mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-08 04:48:35 +08:00
TST: add test for 1d array overwrite
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user