mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-10 20:34:20 +08:00
ENH: add adjustment for datetime64 arrays
BUG: fix adjustment start index
This commit is contained in:
Maya Tydykov
+114
-47
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user