diff --git a/zipline/lib/adjustment.pyx b/zipline/lib/adjustment.pyx index fc133d79..036c5026 100644 --- a/zipline/lib/adjustment.pyx +++ b/zipline/lib/adjustment.pyx @@ -364,7 +364,39 @@ cdef class Float64Overwrite(Float64Adjustment): data[row, col] = value -cdef class Float641DArrayOverwrite: +cdef class ArrayAdjustment(Adjustment): + """ + Base class for ArrayAdjustments. + + Subclasses should inherit and provide a `values` attribute and a `mutate` + method. + """ + def __init__(self, + int64_t first_row, + int64_t last_row, + int64_t first_col, + int64_t last_col): + super(ArrayAdjustment, self).__init__( + first_row=first_row, + last_row=last_row, + first_col=first_col, + last_col=last_col, + ) + + def __repr__(self): + return ( + "%s(first_row=%d, last_row=%d," + " first_col=%d, last_col=%d, values=%s)" % ( + type(self).__name__, + self.first_row, + self.last_row, + self.first_col, + self.last_col, + asarray(self.values), + ) + ) + +cdef class Float641DArrayOverwrite(ArrayAdjustment): """ An adjustment that overwrites subarrays with a value for each subarray. @@ -380,66 +412,101 @@ cdef class Float641DArrayOverwrite: [ 15., 16., 17., 18., 19.], [ 20., 21., 22., 23., 24.]]) >>> adj = Float641DArrayOverwrite( - ... row_starts=np.array([0, 3]), - ... row_ends=np.array([2, 4]), - ... column_starts=np.array([0, 2]), - ... column_ends=np.array([1, 4]), - ... values=np.array([10., 20.]), + ... row_start=0, + ... row_end=3, + ... column_start=0, + ... column_end=0, + ... values=np.array([1, 2, 3, 4]), ) >>> adj.mutate(arr) >>> arr - array([[ 10., 10., 2., 3., 4.], - [ 10., 10., 7., 8., 9.], - [ 10., 10., 12., 13., 14.], - [ 15., 16., 20., 20., 20.], - [ 20., 21., 20., 20., 20.]]) + array([[ 1., 1., 2., 3., 4.], + [ 2., 6., 7., 8., 9.], + [ 3., 11., 12., 13., 14.], + [ 4., 16., 17., 18., 19.], + [ 20., 21., 22., 23., 24.]]) """ cdef: - readonly int64_t[:] row_starts, row_ends, column_starts, column_ends readonly float64_t[:] values def __init__(self, - int64_t[:] row_starts, - int64_t[:] row_ends, - int64_t[:] column_starts, - int64_t[:] column_ends, + int64_t first_row, + int64_t last_row, + int64_t first_col, + int64_t last_col, float64_t[:] values): - assert (len(row_starts) == - len(row_ends) == - len(column_starts) == - len(column_ends)) - for (row_start, row_end) in zip(row_starts, row_ends): - assert row_start <= row_end - for (column_start, column_end) in zip(column_starts, column_ends): - assert column_start <= column_end - - self.row_starts = row_starts - self.row_ends = row_ends - self.column_starts = column_starts - self.column_ends = column_ends + super(Float641DArrayOverwrite, self).__init__( + first_row=first_row, + last_row=last_row, + first_col=first_col, + last_col=last_col, + ) + assert (last_row + 1 - first_row) == len(values) self.values = values cpdef mutate(self, float64_t[:, :] data): cdef Py_ssize_t fill_range, row, col - for fill_range in range(len(self.row_starts)): - for row in range(self.row_starts[fill_range], - self.row_ends[fill_range] + 1): - for col in range(self.column_starts[fill_range], - self.column_ends[fill_range] + 1): - data[row, col] = self.values[fill_range] + cdef float64_t[:] values = self.values + for col in range(self.first_col, self.last_col + 1): + for i, row in enumerate(range(self.first_row, self.last_row + 1)): + data[row, col] = values[i] + + +cdef class Datetime641DArrayOverwrite(ArrayAdjustment): + """ + An adjustment that overwrites subarrays with a value for each subarray. + + Example + ------- + + >>> import numpy as np + >>> arr = np.arange(25, dtype=float).reshape(5, 5) + >>> arr + array([[ 0., 1., 2., 3., 4.], + [ 5., 6., 7., 8., 9.], + [ 10., 11., 12., 13., 14.], + [ 15., 16., 17., 18., 19.], + [ 20., 21., 22., 23., 24.]]) + >>> adj = Datetime641DArrayOverwrite( + ... row_start=0, + ... row_end=3, + ... column_start=0, + ... column_end=0, + ... values=np.array([1, 2, 3, 4]), + ) + >>> adj.mutate(arr) + >>> arr + array([[ 1., 1., 2., 3., 4.], + [ 2., 6., 7., 8., 9.], + [ 3., 11., 12., 13., 14.], + [ 4., 16., 17., 18., 19.], + [ 20., 21., 22., 23., 24.]]) + """ + cdef: + readonly int64_t[:] values + + def __init__(self, + int64_t first_row, + int64_t last_row, + int64_t first_col, + int64_t last_col, + object values): + super(Datetime641DArrayOverwrite, self).__init__( + first_row=first_row, + last_row=last_row, + first_col=first_col, + last_col=last_col, + ) + assert (last_row + 1 - first_row) == len(values) + self.values = asarray([datetime_to_int(value) for value in values]) + + cpdef mutate(self, int64_t[:, :] data): + cdef Py_ssize_t row, col + cdef int64_t[:] values = self.values + for col in range(self.first_col, self.last_col + 1): + for i, row in enumerate(range(self.first_row, self.last_row + 1)): + data[row, col] = values[i] - def __repr__(self): - return ( - "%s(row_starts=%s, row_ends=%s," - " column_starts=%s, column_ends=%s, values=%s)" % ( - type(self).__name__, - asarray(self.row_starts), - asarray(self.row_ends), - asarray(self.column_starts), - asarray(self.column_ends), - asarray(self.values), - ) - ) cdef class Float64Add(Float64Adjustment): """ diff --git a/zipline/pipeline/loaders/quarter_estimates.py b/zipline/pipeline/loaders/quarter_estimates.py index 325e9186..7d9ce0c4 100644 --- a/zipline/pipeline/loaders/quarter_estimates.py +++ b/zipline/pipeline/loaders/quarter_estimates.py @@ -1,11 +1,14 @@ from abc import abstractmethod from collections import defaultdict +from functools import partial import numpy as np +from numpy.ma import asarray import pandas as pd from six import viewvalues -from toolz import groupby +from toolz import groupby, curry from zipline.lib.adjusted_array import AdjustedArray -from zipline.lib.adjustment import Float641DArrayOverwrite +from zipline.lib.adjustment import (Datetime641DArrayOverwrite, + Float641DArrayOverwrite) from zipline.pipeline.common import ( EVENT_DATE_FIELD_NAME, @@ -16,6 +19,7 @@ from zipline.pipeline.common import ( ) from zipline.pipeline.loaders.base import PipelineLoader from zipline.pipeline.loaders.frame import DataFrameLoader +from zipline.utils.numpy_utils import datetime64ns_dtype from zipline.utils.pandas_utils import cross_product from zipline.pipeline.loaders.utils import last_in_date_group, ffill_across_cols @@ -97,42 +101,49 @@ class QuarterEstimatesLoader(PipelineLoader): def load_quarters(self, num_quarters, last, dates): pass - def get_adjustments(self, df, column, mask, assets, - final_releases_per_qtr, dates, raw_events): + def get_adjustments(self, result, col_result, last, + column_name, + column, mask, + assets): 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 + if column.dtype == datetime64ns_dtype: + overwrite = Datetime641DArrayOverwrite + else: + overwrite = Float641DArrayOverwrite + for sid_idx, sid in enumerate(assets): + sid_result = result[result.index.get_level_values( + SID_FIELD_NAME + ) == sid] + sid_result = sid_result.reset_index( + level='shifted_normalized_quarters' + ) # Remove qtrs from index to find shifts + # Figure out where we think quarters are changing. + qtr_shifts = sid_result[ + sid_result['shifted_normalized_quarters'] != + sid_result['shifted_normalized_quarters'].shift(1) ] - # 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)) + # Iterate backwards. No adjustment for 1st quarter. + for row_indexer in list(reversed(qtr_shifts.index))[:-1]: + # We want to write the values for this row's quarter over + # everything that comes before this quarter when we are at + # the date before this quarter starts. + qtr_start_idx = last.index.get_loc(row_indexer[0]) + quarter = qtr_shifts.loc[row_indexer][ + 'shifted_normalized_quarters' + ] + adjustments[qtr_start_idx] = \ + [overwrite(0, + qtr_start_idx - 1, # get index date + sid_idx, + sid_idx, + last[column_name, quarter, + sid][:qtr_start_idx].values) + ] + return AdjustedArray( - df[column.name].values.astype(column.dtype), + col_result.values.astype(column.dtype), mask, - adjustments_from_deltas( - dates, - sparse_output[TS_FIELD_NAME].values, - column_idx, - column.name, - asset_idx, - sparse_deltas, - ), + dict(adjustments), column.missing_value, ) @@ -150,8 +161,6 @@ class QuarterEstimatesLoader(PipelineLoader): date_values[SIMULTATION_DATES] = date_values[ SIMULTATION_DATES ].astype('datetime64[ns]') - 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], @@ -166,50 +175,51 @@ class QuarterEstimatesLoader(PipelineLoader): last = last_in_date_group(self.estimates, True, dates, assets, extra_groupers=[ - 'normalized_quarters']).reset_index() + 'normalized_quarters']) # Forward fill values for each quarter. ffill_across_cols(last, columns) - stacked = last.stack(1).stack(1).reset_index() + stacked = last.stack(1).stack(1) - result = self.load_quarters(num_quarters, - stacked, dates) + result = self.load_quarters(num_quarters, stacked) 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( - c, - result.pivot(index=SIMULTATION_DATES, - columns=SID_FIELD_NAME, - values=column_name), - adjustments=adjusted_array - ) - out[c] = loader.load_adjusted_array([c], - dates, - assets, - mask)[c] + col_result = result[ + column_name + ].reset_index(1, drop=True).unstack(1).reindex(dates) + adjusted_array = self.get_adjustments(result, + col_result, + last, + column_name, + c, + mask, + assets) + out[c] = adjusted_array return out class NextQuartersEstimatesLoader(QuarterEstimatesLoader): - def load_quarters(self, num_quarters, stacked, dates): + def load_quarters(self, num_quarters, stacked): # 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 + stacked = stacked.sort(EVENT_DATE_FIELD_NAME) + next_releases = stacked.loc[ + stacked[EVENT_DATE_FIELD_NAME] >= stacked.index.get_level_values( + 0 + )].groupby(level=[0, 2]).nth(0) + next_releases[ + 'shifted_normalized_quarters' + ] = next_releases.index.get_level_values( + 'normalized_quarters' + ) + (num_quarters - 1) + next_releases = next_releases.set_index([ + next_releases.index.get_level_values(0), # dates + 'shifted_normalized_quarters', + next_releases.index.get_level_values(2) # sids + ]) + return stacked.loc[next_releases.index] class PreviousQuartersEstimatesLoader(QuarterEstimatesLoader):