TST: add test for 1d array overwrite

This commit is contained in:
Maya Tydykov
2016-08-18 14:19:48 -04:00
parent e53d7fc9b2
commit 2975f9b2fd
5 changed files with 261 additions and 120 deletions
+101
View File
@@ -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
+13 -68
View File
@@ -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
+76 -51
View File
@@ -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
+71
View File
@@ -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)