diff --git a/tests/pipeline/test_adjusted_array.py b/tests/pipeline/test_adjusted_array.py index cea6f090..80a1dd97 100644 --- a/tests/pipeline/test_adjusted_array.py +++ b/tests/pipeline/test_adjusted_array.py @@ -22,6 +22,7 @@ from zipline.lib.adjustment import ( Datetime64Overwrite, Float64Multiply, Float64Overwrite, + Float641DArrayOverwrite, ObjectOverwrite, ) from zipline.lib.adjusted_array import AdjustedArray, NOMASK @@ -304,6 +305,105 @@ def _gen_overwrite_adjustment_cases(name, ) +def _gen_overwrite_1d_array_adjustment_case(): + """ + Generate test cases for overwrite adjustments. + + The algorithm used here is the same as the one used above for + multiplicative adjustments. The only difference is the semantics of how + the adjustments are expected to modify the arrays. + + This is parameterized on `make_input` and `make_expected_output` functions, + which take 2-D lists of values and transform them into desired input/output + arrays. We do this so that we can easily test both vanilla numpy ndarrays + and our own LabelArray class for strings. + """ + + adjustments = {} + buffer_as_of = [None] * 6 + baseline = as_dtype(float64_dtype, [[2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2]]) + + buffer_as_of[0] = as_dtype(float64_dtype, [[2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2]]) + + # Note that row indices are inclusive! + adjustments[1] = [ + Float641DArrayOverwrite(array([0]), + array([0]), + array([0]), + array([0]), + as_dtype(float64_dtype, array([1]))) + ] + buffer_as_of[1] = as_dtype(float64_dtype, [[1, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2]]) + + # No adjustment at index 2. + buffer_as_of[2] = buffer_as_of[1] + + adjustments[3] = [ + Float641DArrayOverwrite(array([0, 2, 1]), + array([1, 2, 2]), + array([0, 0, 1]), + array([0, 0, 1]), + as_dtype(float64_dtype, array([4, 1, 3]))) + ] + buffer_as_of[3] = as_dtype(float64_dtype, [[4, 2, 2], + [4, 3, 2], + [1, 3, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2]]) + + adjustments[4] = [ + Float641DArrayOverwrite(array([0]), + array([3]), + array([2]), + array([2]), + as_dtype(float64_dtype, array([5]))) + ] + buffer_as_of[4] = as_dtype(float64_dtype, [[4, 2, 5], + [4, 3, 5], + [1, 3, 5], + [2, 2, 5], + [2, 2, 2], + [2, 2, 2]]) + + adjustments[5] = [ + Float641DArrayOverwrite(array([0, 2]), + array([4, 2]), + array([1, 2]), + array([1, 2]), + as_dtype(float64_dtype, array([6, 7]))), + ] + buffer_as_of[5] = as_dtype(float64_dtype, [[4, 6, 5], + [4, 6, 5], + [1, 6, 7], + [2, 6, 5], + [2, 6, 2], + [2, 2, 2]]) + + return _gen_expectations( + baseline, + default_missing_value_for_dtype(float64_dtype), + adjustments, + buffer_as_of, + nrows=6, + ) + + def _gen_expectations(baseline, missing_value, adjustments, @@ -442,6 +542,7 @@ class AdjustedArrayTestCase(TestCase): datetime64ns_dtype, ), ), + _gen_overwrite_1d_array_adjustment_case(), # There are six cases here: # Using np.bytes/np.unicode/object arrays as inputs. # Passing np.bytes/np.unicode/object arrays to LabelArray, diff --git a/tests/pipeline/test_quarters_estimates.py b/tests/pipeline/test_quarters_estimates.py index f343bfb8..1f82d161 100644 --- a/tests/pipeline/test_quarters_estimates.py +++ b/tests/pipeline/test_quarters_estimates.py @@ -20,7 +20,6 @@ from zipline.pipeline.loaders.quarter_estimates import ( NextQuartersEstimatesLoader, PreviousQuartersEstimatesLoader ) -from zipline.pipeline.loaders.quarter_estimates import shift_quarters from zipline.testing import ZiplineTestCase from zipline.testing.fixtures import WithAssetFinder, WithTradingSessions from zipline.testing.predicates import assert_equal diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 01e1e659..c076029d 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -175,9 +175,10 @@ from zipline.pipeline.common import ( from zipline.pipeline.data.dataset import DataSet, Column from zipline.pipeline.loaders.utils import ( check_data_query_args, + last_in_date_group, normalize_data_query_bounds, normalize_timestamp_to_query_time, -) + ffill_across_cols) from zipline.pipeline.sentinels import NotSpecified from zipline.lib.adjusted_array import AdjustedArray, can_represent_dtype from zipline.lib.adjustment import Float64Overwrite @@ -869,9 +870,9 @@ def adjustments_from_deltas_with_sids(dense_dates, Parameters ---------- - dates : pd.DatetimeIndex - The dates requested by the loader. dense_dates : pd.DatetimeIndex + The dates requested by the loader. + sparse_dates : pd.DatetimeIndex The dates that were in the raw data. column_idx : int The index of the column in the dataset. @@ -1091,71 +1092,15 @@ class BlazeLoader(dict): ) sparse_output.drop(AD_FIELD_NAME, axis=1, inplace=True) - def last_in_date_group(df, reindex, have_sids=have_sids): - idx = dates[dates.searchsorted( - df[TS_FIELD_NAME].values.astype('datetime64[D]') - )] - if have_sids: - idx = [idx, SID_FIELD_NAME] - - last_in_group = df.drop(TS_FIELD_NAME, axis=1).groupby( - idx, - sort=False, - ).last() - - if have_sids: - last_in_group = last_in_group.unstack() - - if reindex: - if have_sids: - cols = last_in_group.columns - last_in_group = last_in_group.reindex( - index=dates, - columns=pd.MultiIndex.from_product( - (cols.levels[0], assets), - names=cols.names, - ), - ) - else: - last_in_group = last_in_group.reindex(dates) - - return last_in_group - - sparse_deltas = last_in_date_group(non_novel_deltas, reindex=False) - dense_output = last_in_date_group(sparse_output, reindex=True) - dense_output.ffill(inplace=True) - - # Fill in missing values specified by each column. This is made - # significantly more complex by the fact that we need to work around - # two pandas issues: - - # 1) When we have sids, if there are no records for a given sid for any - # dates, pandas will generate a column full of NaNs for that sid. - # This means that some of the columns in `dense_output` are now - # float instead of the intended dtype, so we have to coerce back to - # our expected type and convert NaNs into the desired missing value. - - # 2) DataFrame.ffill assumes that receiving None as a fill-value means - # that no value was passed. Consequently, there's no way to tell - # pandas to replace NaNs in an object column with None using fillna, - # so we have to roll our own instead using df.where. - for column in columns: - # Special logic for strings since `fillna` doesn't work if the - # missing value is `None`. - if column.dtype == categorical_dtype: - dense_output[column.name] = dense_output[ - column.name - ].where(pd.notnull(dense_output[column.name]), - column.missing_value) - else: - # We need to execute `fillna` before `astype` in case the - # column contains NaNs and needs to be cast to bool or int. - # This is so that the NaNs are replaced first, since pandas - # can't convert NaNs for those types. - dense_output[column.name] = dense_output[ - column.name - ].fillna(column.missing_value).astype(column.dtype) - + sparse_deltas = last_in_date_group(non_novel_deltas, + dates, + assets, + reindex=False) + dense_output = last_in_date_group(sparse_output, + dates, + assets, + reindex=True) + ffill_across_cols(dense_output, columns) if have_sids: adjustments_from_deltas = adjustments_from_deltas_with_sids column_view = identity diff --git a/zipline/pipeline/loaders/quarter_estimates.py b/zipline/pipeline/loaders/quarter_estimates.py index 7837e57d..325e9186 100644 --- a/zipline/pipeline/loaders/quarter_estimates.py +++ b/zipline/pipeline/loaders/quarter_estimates.py @@ -1,7 +1,11 @@ from abc import abstractmethod +from collections import defaultdict +import numpy as np import pandas as pd from six import viewvalues from toolz import groupby +from zipline.lib.adjusted_array import AdjustedArray +from zipline.lib.adjustment import Float641DArrayOverwrite from zipline.pipeline.common import ( EVENT_DATE_FIELD_NAME, @@ -13,6 +17,7 @@ from zipline.pipeline.common import ( from zipline.pipeline.loaders.base import PipelineLoader from zipline.pipeline.loaders.frame import DataFrameLoader from zipline.utils.pandas_utils import cross_product +from zipline.pipeline.loaders.utils import last_in_date_group, ffill_across_cols NEXT_FISCAL_QUARTER = 'next_fiscal_quarter' NEXT_FISCAL_YEAR = 'next_fiscal_year' @@ -31,10 +36,6 @@ def split_normalized_quarters(normalized_quarters): return years, quarters + 1 -def shift_quarters(by, years, quarters): - return split_normalized_quarters(normalize_quarters(years, quarters) + by) - - def required_estimates_fields(columns): """ Compute the set of resource columns required to serve @@ -93,15 +94,54 @@ class QuarterEstimatesLoader(PipelineLoader): self.base_column_name_map = base_column_name_map @abstractmethod - def load_quarters(self, num_quarters, dates_sids, final_releases_per_qtr): + def load_quarters(self, num_quarters, last, dates): pass + def get_adjustments(self, df, column, mask, assets, + final_releases_per_qtr, dates, raw_events): + adjustments = defaultdict(list) + for idx, sid in enumerate(assets): + # Get the releases for a particular sid + sid_data = final_releases_per_qtr[final_releases_per_qtr[ + SID_FIELD_NAME] == sid + ] + # Get the release dates for this sid - these are the quarter + # boundaries + qtr_boundaries, years, qtrs = sid_data[[ + EVENT_DATE_FIELD_NAME, + FISCAL_YEAR_FIELD_NAME, + FISCAL_QUARTER_FIELD_NAME + ]].unique() + next_qtr_starts = dates.searchsorted(qtr_boundaries, sid='right') + for idx, start in enumerate(next_qtr_starts): + # Here we need to take the new quarter and, for all dates in + # previous quarters, apply adjustments that use this + # quarter's values for those previous dates. + adjustments[start].extend(Float641DArrayOverwrite(first_row, + last_row, + idx, + idx, + value)) + return AdjustedArray( + df[column.name].values.astype(column.dtype), + mask, + adjustments_from_deltas( + dates, + sparse_output[TS_FIELD_NAME].values, + column_idx, + column.name, + asset_idx, + sparse_deltas, + ), + column.missing_value, + ) + def load_adjusted_array(self, columns, dates, assets, mask): # TODO: how can we enforce that datasets have the num_quarters # attribute, given that they're created dynamically? groups = groupby(lambda x: x.dataset.num_quarters, columns) groups_columns = dict(groups) - if (pd.Series(groups_columns.keys()) < 0).any(): + if (pd.Series(groups_columns) < 0).any(): raise ValueError("Must pass a number of quarters >= 0") out = {} date_values = pd.DataFrame({SIMULTATION_DATES: dates}) @@ -110,34 +150,36 @@ class QuarterEstimatesLoader(PipelineLoader): date_values[SIMULTATION_DATES] = date_values[ SIMULTATION_DATES ].astype('datetime64[ns]') - estimates_all_dates = cross_product(date_values, self.estimates) asset_df = pd.DataFrame({SID_FIELD_NAME: assets}) dates_sids = cross_product(date_values, asset_df) + self.estimates['normalized_quarters'] = normalize_quarters( + self.estimates[FISCAL_YEAR_FIELD_NAME], + self.estimates[FISCAL_QUARTER_FIELD_NAME], + ).astype(float) for num_quarters, columns in groups_columns.iteritems(): name_map = {c: self.base_column_name_map[ getattr(c.dataset.__base__, c.name) ] for c in columns} - - # First, determine which estimates we would have known about on - # each date. Then, Sort by timestamp and group to find the latest - # estimate for each quarter. - final_releases_per_qtr = estimates_all_dates[ - estimates_all_dates[TS_FIELD_NAME] <= - estimates_all_dates.dates - ].sort([TS_FIELD_NAME]).groupby( - [SIMULTATION_DATES, - SID_FIELD_NAME, - FISCAL_YEAR_FIELD_NAME, - FISCAL_QUARTER_FIELD_NAME] - ).nth(-1).reset_index() + # Determine the last piece of information we know for each column + # on each date in the index. + last = last_in_date_group(self.estimates, True, dates, + assets, + extra_groupers=[ + 'normalized_quarters']).reset_index() + # Forward fill values for each quarter. + ffill_across_cols(last, columns) + stacked = last.stack(1).stack(1).reset_index() result = self.load_quarters(num_quarters, - dates_sids, - final_releases_per_qtr) + stacked, dates) for c in columns: column_name = name_map[c] + pivoted = result.pivot(index=SIMULTATION_DATES, + columns=SID_FIELD_NAME, + values=column_name) + adjusted_array = self.get_adjustments(pivoted, c, mask, assets) # Pivot to get a DataFrame with dates as the index and # sids as the columns. loader = DataFrameLoader( @@ -145,7 +187,7 @@ class QuarterEstimatesLoader(PipelineLoader): result.pivot(index=SIMULTATION_DATES, columns=SID_FIELD_NAME, values=column_name), - adjustments=None + adjustments=adjusted_array ) out[c] = loader.load_adjusted_array([c], dates, @@ -156,34 +198,17 @@ class QuarterEstimatesLoader(PipelineLoader): class NextQuartersEstimatesLoader(QuarterEstimatesLoader): - def load_quarters(self, num_quarters, dates_sids, final_releases_per_qtr): - # Filter for releases that are on or after each simulation date. - eligible_next_releases = final_releases_per_qtr[ - final_releases_per_qtr[EVENT_DATE_FIELD_NAME] >= - final_releases_per_qtr[SIMULTATION_DATES] - ] - # For each sid, get the upcoming release. - eligible_next_releases.sort(EVENT_DATE_FIELD_NAME) - next_releases = eligible_next_releases.groupby( - [SIMULTATION_DATES, SID_FIELD_NAME] - ).nth(0).reset_index() # We use nth here to avoid forward filling - # NaNs, which `first()` will do. - next_releases = next_releases.rename( - columns={FISCAL_YEAR_FIELD_NAME: NEXT_FISCAL_YEAR, - FISCAL_QUARTER_FIELD_NAME: NEXT_FISCAL_QUARTER} - ) - # The next fiscal quarter is already our starting point, - # so we should offset `num_quarters` by 1. - (next_releases[FISCAL_YEAR_FIELD_NAME], - next_releases[FISCAL_QUARTER_FIELD_NAME]) = shift_quarters( - (num_quarters - 1), - next_releases[NEXT_FISCAL_YEAR], - next_releases[NEXT_FISCAL_QUARTER], - ) - # Do a left merge to get values for each date. - result = dates_sids.merge(next_releases, - on=([SIMULTATION_DATES, SID_FIELD_NAME]), - how='left') + def load_quarters(self, num_quarters, stacked, dates): + # Filter for releases that are on or after each simulation date and + # determine the next quarter by picking out the upcoming release for + # each date in the index. + event_date_idxs = dates.searchsorted(pd.to_datetime(stacked[EVENT_DATE_FIELD_NAME]).values) + next_releases = stacked.loc[event_date_idxs >= stacked['level_0']].groupby(['level_0', 'sid']).nth(0) + + + next_releases['shifted_normalized_quarters'] = next_releases[ + 'normalized_quarters'].convert_objects(convert_numeric=True) + (num_quarters - 1) + return result diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py index 77a9f447..28e0c90d 100644 --- a/zipline/pipeline/loaders/utils.py +++ b/zipline/pipeline/loaders/utils.py @@ -2,6 +2,8 @@ import datetime import numpy as np import pandas as pd +from zipline.pipeline.common import TS_FIELD_NAME, SID_FIELD_NAME +from zipline.utils.numpy_utils import categorical_dtype from zipline.utils.pandas_utils import mask_between_time @@ -272,3 +274,72 @@ def check_data_query_args(data_query_time, data_query_tz): data_query_tz, ), ) + + +def last_in_date_group(df, reindex, dates, assets, have_sids=True, + extra_groupers=[]): + idx = dates[dates.searchsorted( + df[TS_FIELD_NAME].values.astype('datetime64[D]') + )] + if have_sids: + idx = [idx, SID_FIELD_NAME] + extra_groupers + + last_in_group = df.drop(TS_FIELD_NAME, axis=1).groupby( + idx, + sort=False, + ).last() + + # For the number of things that we're grouping by (except TS), unstack + # the df + for _ in range(len(idx) - 1): + last_in_group = last_in_group.unstack() + + if reindex: + if have_sids: + cols = last_in_group.columns + last_in_group = last_in_group.reindex( + index=dates, + columns=pd.MultiIndex.from_product( + tuple(cols.levels[0:len(extra_groupers) + 1]) + (assets,), + names=cols.names, + ), + ) + else: + last_in_group = last_in_group.reindex(dates) + + return last_in_group + + +def ffill_across_cols(df, columns): + df.ffill(inplace=True) + + # Fill in missing values specified by each column. This is made + # significantly more complex by the fact that we need to work around + # two pandas issues: + + # 1) When we have sids, if there are no records for a given sid for any + # dates, pandas will generate a column full of NaNs for that sid. + # This means that some of the columns in `dense_output` are now + # float instead of the intended dtype, so we have to coerce back to + # our expected type and convert NaNs into the desired missing value. + + # 2) DataFrame.ffill assumes that receiving None as a fill-value means + # that no value was passed. Consequently, there's no way to tell + # pandas to replace NaNs in an object column with None using fillna, + # so we have to roll our own instead using df.where. + for column in columns: + # Special logic for strings since `fillna` doesn't work if the + # missing value is `None`. + if column.dtype == categorical_dtype: + df[column.name] = df[ + column.name + ].where(pd.notnull(df[column.name]), + column.missing_value) + else: + # We need to execute `fillna` before `astype` in case the + # column contains NaNs and needs to be cast to bool or int. + # This is so that the NaNs are replaced first, since pandas + # can't convert NaNs for those types. + df[column.name] = df[ + column.name + ].fillna(column.missing_value).astype(column.dtype)