def expected_adjustments(self, start_date, end_date, tables,
adjustment_type):
price_adjustments = {}
volume_adjustments = {}
should_include_price_adjustments = (adjustment_type == 'all'
or adjustment_type == 'price')
should_include_volume_adjustments = (adjustment_type == 'all'
or adjustment_type == 'volume')
query_days = self.calendar_days_between(start_date, end_date)
start_loc = query_days.get_loc(start_date)
for table in tables:
for eff_date_secs, ratio, sid in table.itertuples(index=False):
eff_date = Timestamp(eff_date_secs, unit='s', tz='UTC')
# Ignore adjustments outside the query bounds.
if not (start_date <= eff_date <= end_date):
continue
eff_date_loc = query_days.get_loc(eff_date)
delta = eff_date_loc - start_loc
# Pricing adjustments should be applied on the date
# corresponding to the effective date of the input data. They
# should affect all rows **before** the effective date.
if should_include_price_adjustments:
price_adjustments.setdefault(delta, []).append(
Float64Multiply(
first_row=0,
last_row=delta,
first_col=sid - 1,
last_col=sid - 1,
value=ratio,
))
# Volume is *inversely* affected by *splits only*.
if table is SPLITS and should_include_volume_adjustments:
volume_adjustments.setdefault(delta, []).append(
Float64Multiply(
first_row=0,
last_row=delta,
first_col=sid - 1,
last_col=sid - 1,
value=1.0 / ratio,
))
output = {}
if should_include_price_adjustments:
output['price_adjustments'] = price_adjustments
if should_include_volume_adjustments:
output['volume_adjustments'] = volume_adjustments
return output
def test_inspect(self):
data = arange(15, dtype=float).reshape(5, 3)
adj_array = AdjustedArray(
data,
{4: [Float64Multiply(2, 3, 0, 0, 4.0)]},
float('nan'),
)
expected = dedent(
"""\
Adjusted Array (float64):
Data:
array([[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.],
[ 9., 10., 11.],
[ 12., 13., 14.]])
Adjustments:
{4: [Float64Multiply(first_row=2, last_row=3, first_col=0, \
last_col=0, value=4.000000)]}
"""
)
got = adj_array.inspect()
self.assertEqual(expected, got)
def test_inspect(self):
data = np.arange(15, dtype=float).reshape(5, 3)
adj_array = AdjustedArray(
data,
{4: [Float64Multiply(2, 3, 0, 0, 4.0)]},
float("nan"),
)
# TODO: CHECK WHY DO I NEED TO FIX THE INDENT IN THE EXPECTED?
expected = dedent(
"""\
Adjusted Array (float64):
Data:
array([[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.],
[ 9., 10., 11.],
[12., 13., 14.]])
Adjustments:
{4: [Float64Multiply(first_row=2, last_row=3, first_col=0, \
last_col=0, value=4.000000)]}
"""
)
got = adj_array.inspect()
assert expected == got
def test_copy(self):
data = arange(5 * 3, dtype='f8').reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float('nan'))
traverse_copy = adjusted_array.copy()
clean_copy = adjusted_array.copy()
a_it = adjusted_array.traverse(2, copy=False)
b_it = traverse_copy.traverse(2, copy=False)
for a, b in zip(a_it, b_it):
assert_equal(a, b)
with self.assertRaises(ValueError) as e:
adjusted_array.copy()
assert_equal(
str(e.exception),
'cannot copy invalidated AdjustedArray',
)
# the clean copy should have the original data even though the
# original adjusted array has it's data mutated in place
assert_equal(clean_copy.data, original_data)
assert_equal(adjusted_array.data, original_data * 2)
def expected_adjustments(self, start_date, end_date):
price_adjustments = {}
volume_adjustments = {}
query_days = self.calendar_days_between(start_date, end_date)
start_loc = query_days.get_loc(start_date)
for table in SPLITS, MERGERS, DIVIDENDS_EXPECTED:
for eff_date_secs, ratio, sid in table.itertuples(index=False):
eff_date = Timestamp(eff_date_secs, unit='s', tz='UTC')
# Ignore adjustments outside the query bounds.
if not (start_date <= eff_date <= end_date):
continue
eff_date_loc = query_days.get_loc(eff_date)
delta = eff_date_loc - start_loc
# Pricing adjustments should be applied on the date
# corresponding to the effective date of the input data. They
# should affect all rows **before** the effective date.
price_adjustments.setdefault(delta, []).append(
Float64Multiply(
first_row=0,
last_row=delta,
first_col=sid - 1,
last_col=sid - 1,
value=ratio,
)
)
# Volume is *inversely* affected by *splits only*.
if table is SPLITS:
volume_adjustments.setdefault(delta, []).append(
Float64Multiply(
first_row=0,
last_row=delta,
first_col=sid - 1,
last_col=sid - 1,
value=1.0 / ratio,
)
)
return price_adjustments, volume_adjustments
def expected_adjustments(self, start_date, end_date):
price_adjustments = {}
volume_adjustments = {}
query_days = self.calendar_days_between(start_date, end_date)
start_loc = query_days.get_loc(start_date)
for table in SPLITS, MERGERS, DIVIDENDS:
for eff_date_secs, ratio, sid in table.itertuples(index=False):
eff_date = Timestamp(eff_date_secs, unit='s', tz='UTC')
# The boundary conditions here are subtle. An adjustment with
# an effective date equal to the query start can't have an
# effect because adjustments only the array for dates strictly
# less than the adjustment effective date.
if not (start_date < eff_date <= end_date):
continue
eff_date_loc = query_days.get_loc(eff_date)
delta = eff_date_loc - start_loc
# Pricing adjusments should be applied on the date
# corresponding to the effective date of the input data. They
# should affect all rows **before** the effective date.
price_adjustments.setdefault(delta, []).append(
Float64Multiply(
first_row=0,
last_row=delta - 1,
col=sid - 1,
value=ratio,
))
# Volume is *inversely* affected by *splits only*.
if table is SPLITS:
volume_adjustments.setdefault(delta, []).append(
Float64Multiply(
first_row=0,
last_row=delta - 1,
col=sid - 1,
value=1.0 / ratio,
))
return price_adjustments, volume_adjustments
def split_adjustment(sid, volume):
"""The splits occur at index 252 // 2 with a ratio of (sid + 1):1
"""
idx = 252 // 2
return {
idx: [Float64Multiply(
first_row=0,
last_row=idx,
first_col=sid,
last_col=sid,
value=(identity if volume else op.truediv(1))(sid + 2),
)],
}
def test_traverse_invalidating(self):
data = np.arange(5 * 3, dtype="f8").reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float("nan"))
for _ in adjusted_array.traverse(1, copy=False):
pass
assert_equal(data, original_data * 2)
err_msg = "cannot traverse invalidated AdjustedArray"
with pytest.raises(ValueError, match=err_msg):
adjusted_array.traverse(1)
def test_traverse_invalidating(self):
data = arange(5 * 3, dtype='f8').reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float('nan'))
for _ in adjusted_array.traverse(1, copy=False):
pass
assert_equal(data, original_data * 2)
with self.assertRaises(ValueError) as e:
adjusted_array.traverse(1)
assert_equal(
str(e.exception),
'cannot traverse invalidated AdjustedArray',
)
def test_copy(self):
data = np.arange(5 * 3, dtype="f8").reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float("nan"))
traverse_copy = adjusted_array.copy()
clean_copy = adjusted_array.copy()
a_it = adjusted_array.traverse(2, copy=False)
b_it = traverse_copy.traverse(2, copy=False)
for a, b in zip(a_it, b_it):
assert_equal(a, b)
err_msg = "cannot copy invalidated AdjustedArray"
with pytest.raises(ValueError, match=err_msg):
adjusted_array.copy()
# the clean copy should have the original data even though the
# original adjusted array has it's data mutated in place
assert_equal(clean_copy.data, original_data)
assert_equal(adjusted_array.data, original_data * 2)
def dividend_adjustment(sid, which):
"""The dividends occur at indices 252 // 4 and 3 * 252 / 4
with a cash amount of sid + 1 / 10 and sid + 2 / 10
"""
if which == 'first':
idx = 252 // 4
else:
idx = 3 * 252 // 4
return {
idx: [
Float64Multiply(
first_row=0,
last_row=idx,
first_col=sid,
last_col=sid,
value=float(1 - ((sid + 1 +
(which == 'second')) / 10) /
(idx - 1 + sid * 10000 + 2000)),
)
],
}
class AdjustedArrayTestCase(TestCase):
def test_traverse_invalidating(self):
data = arange(5 * 3, dtype='f8').reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float('nan'))
for _ in adjusted_array.traverse(1, copy=False):
pass
assert_equal(data, original_data * 2)
with self.assertRaises(ValueError) as e:
adjusted_array.traverse(1)
assert_equal(
str(e.exception),
'cannot traverse invalidated AdjustedArray',
)
def test_copy(self):
data = arange(5 * 3, dtype='f8').reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float('nan'))
traverse_copy = adjusted_array.copy()
clean_copy = adjusted_array.copy()
a_it = adjusted_array.traverse(2, copy=False)
b_it = traverse_copy.traverse(2, copy=False)
for a, b in zip(a_it, b_it):
assert_equal(a, b)
with self.assertRaises(ValueError) as e:
adjusted_array.copy()
assert_equal(
str(e.exception),
'cannot copy invalidated AdjustedArray',
)
# the clean copy should have the original data even though the
# original adjusted array has it's data mutated in place
assert_equal(clean_copy.data, original_data)
assert_equal(adjusted_array.data, original_data * 2)
@parameterized.expand(
chain(
_gen_unadjusted_cases(
'float',
make_input=as_dtype(float64_dtype),
make_expected_output=as_dtype(float64_dtype),
missing_value=default_missing_value_for_dtype(float64_dtype),
),
_gen_unadjusted_cases(
'datetime',
make_input=as_dtype(datetime64ns_dtype),
make_expected_output=as_dtype(datetime64ns_dtype),
missing_value=default_missing_value_for_dtype(
datetime64ns_dtype),
),
# Test passing an array of strings to AdjustedArray.
_gen_unadjusted_cases(
'bytes_ndarray',
make_input=as_dtype(bytes_dtype),
make_expected_output=as_labelarray(bytes_dtype, b''),
missing_value=b'',
),
_gen_unadjusted_cases(
'unicode_ndarray',
make_input=as_dtype(unicode_dtype),
make_expected_output=as_labelarray(unicode_dtype, u''),
missing_value=u'',
),
_gen_unadjusted_cases(
'object_ndarray',
make_input=lambda a: a.astype(unicode).astype(object),
make_expected_output=as_labelarray(unicode_dtype, u''),
missing_value='',
),
# Test passing a LabelArray directly to AdjustedArray.
_gen_unadjusted_cases(
'bytes_labelarray',
make_input=as_labelarray(bytes_dtype, b''),
make_expected_output=as_labelarray(bytes_dtype, b''),
missing_value=b'',
),
_gen_unadjusted_cases(
'unicode_labelarray',
make_input=as_labelarray(unicode_dtype, None),
make_expected_output=as_labelarray(unicode_dtype, None),
missing_value=u'',
),
_gen_unadjusted_cases(
'object_labelarray',
make_input=(lambda a: LabelArray(
a.astype(unicode).astype(object), u'')),
make_expected_output=as_labelarray(unicode_dtype, ''),
missing_value='',
),
))
def test_no_adjustments(self, name, data, lookback, adjustments,
missing_value, perspective_offset,
expected_output):
array = AdjustedArray(data, adjustments, missing_value)
for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable.
in_out = zip(array.traverse(lookback), expected_output)
for yielded, expected_yield in in_out:
check_arrays(yielded, expected_yield)
@parameterized.expand(_gen_multiplicative_adjustment_cases(float64_dtype))
def test_multiplicative_adjustments(self, name, data, lookback,
adjustments, missing_value,
perspective_offset, expected):
array = AdjustedArray(data, adjustments, missing_value)
for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable.
window_iter = array.traverse(
lookback,
perspective_offset=perspective_offset,
)
for yielded, expected_yield in zip_longest(window_iter, expected):
check_arrays(yielded, expected_yield)
@parameterized.expand(
chain(
_gen_overwrite_adjustment_cases(bool_dtype),
_gen_overwrite_adjustment_cases(int64_dtype),
_gen_overwrite_adjustment_cases(float64_dtype),
_gen_overwrite_adjustment_cases(datetime64ns_dtype),
_gen_overwrite_1d_array_adjustment_case(float64_dtype),
_gen_overwrite_1d_array_adjustment_case(datetime64ns_dtype),
_gen_overwrite_1d_array_adjustment_case(bool_dtype),
# There are six cases here:
# Using np.bytes/np.unicode/object arrays as inputs.
# Passing np.bytes/np.unicode/object arrays to LabelArray,
# and using those as input.
#
# The outputs should always be LabelArrays.
_gen_unadjusted_cases(
'bytes_ndarray',
make_input=as_dtype(bytes_dtype),
make_expected_output=as_labelarray(bytes_dtype, b''),
missing_value=b'',
),
_gen_unadjusted_cases(
'unicode_ndarray',
make_input=as_dtype(unicode_dtype),
make_expected_output=as_labelarray(unicode_dtype, u''),
missing_value=u'',
),
_gen_unadjusted_cases(
'object_ndarray',
make_input=lambda a: a.astype(unicode).astype(object),
make_expected_output=as_labelarray(unicode_dtype, u''),
missing_value=u'',
),
_gen_unadjusted_cases(
'bytes_labelarray',
make_input=as_labelarray(bytes_dtype, b''),
make_expected_output=as_labelarray(bytes_dtype, b''),
missing_value=b'',
),
_gen_unadjusted_cases(
'unicode_labelarray',
make_input=as_labelarray(unicode_dtype, u''),
make_expected_output=as_labelarray(unicode_dtype, u''),
missing_value=u'',
),
_gen_unadjusted_cases(
'object_labelarray',
make_input=(lambda a: LabelArray(
a.astype(unicode).astype(object),
None,
)),
make_expected_output=as_labelarray(unicode_dtype, u''),
missing_value=None,
),
))
def test_overwrite_adjustment_cases(self, name, baseline, lookback,
adjustments, missing_value,
perspective_offset, expected):
array = AdjustedArray(baseline, adjustments, missing_value)
for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable.
window_iter = array.traverse(
lookback,
perspective_offset=perspective_offset,
)
for yielded, expected_yield in zip_longest(window_iter, expected):
check_arrays(yielded, expected_yield)
def test_object1darrayoverwrite(self):
pairs = [u + l for u, l in product(ascii_uppercase, ascii_lowercase)]
categories = pairs + ['~' + c for c in pairs]
baseline = LabelArray(
array([[''.join((r, c)) for c in 'abc'] for r in ascii_uppercase]),
None,
categories,
)
full_expected = baseline.copy()
def flip(cs):
if cs is None:
return None
if cs[0] != '~':
return '~' + cs
return cs
def make_overwrite(fr, lr, fc, lc):
fr, lr, fc, lc = map(ord, (fr, lr, fc, lc))
fr -= ord('A')
lr -= ord('A')
fc -= ord('a')
lc -= ord('a')
return Object1DArrayOverwrite(
fr,
lr,
fc,
lc,
baseline[fr:lr + 1, fc].map(flip),
)
overwrites = {
3: [make_overwrite('A', 'B', 'a', 'a')],
4: [make_overwrite('A', 'C', 'b', 'c')],
5: [make_overwrite('D', 'D', 'a', 'b')],
}
it = AdjustedArray(baseline, overwrites, None).traverse(3)
window = next(it)
expected = full_expected[:3]
check_arrays(window, expected)
window = next(it)
full_expected[0:2, 0] = LabelArray(['~Aa', '~Ba'], None)
expected = full_expected[1:4]
check_arrays(window, expected)
window = next(it)
full_expected[0:3, 1:3] = LabelArray(
[['~Ab', '~Ac'], ['~Bb', '~Bc'], ['~Cb', '~Cb']], None)
expected = full_expected[2:5]
check_arrays(window, expected)
window = next(it)
full_expected[3, :2] = '~Da'
expected = full_expected[3:6]
check_arrays(window, expected)
def test_invalid_lookback(self):
data = arange(30, dtype=float).reshape(6, 5)
adj_array = AdjustedArray(data, {}, float('nan'))
with self.assertRaises(WindowLengthTooLong):
adj_array.traverse(7)
with self.assertRaises(WindowLengthNotPositive):
adj_array.traverse(0)
with self.assertRaises(WindowLengthNotPositive):
adj_array.traverse(-1)
def test_array_views_arent_writable(self):
data = arange(30, dtype=float).reshape(6, 5)
adj_array = AdjustedArray(data, {}, float('nan'))
for frame in adj_array.traverse(3):
with self.assertRaises(ValueError):
frame[0, 0] = 5.0
def test_inspect(self):
data = arange(15, dtype=float).reshape(5, 3)
adj_array = AdjustedArray(
data,
{4: [Float64Multiply(2, 3, 0, 0, 4.0)]},
float('nan'),
)
expected = dedent("""\
Adjusted Array (float64):
Data:
array([[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.],
[ 9., 10., 11.],
[ 12., 13., 14.]])
Adjustments:
{4: [Float64Multiply(first_row=2, last_row=3, first_col=0, \
last_col=0, value=4.000000)]}
""")
got = adj_array.inspect()
self.assertEqual(expected, got)
def test_update_labels(self):
data = array([
['aaa', 'bbb', 'ccc'],
['ddd', 'eee', 'fff'],
['ggg', 'hhh', 'iii'],
['jjj', 'kkk', 'lll'],
['mmm', 'nnn', 'ooo'],
])
label_array = LabelArray(data, missing_value='')
adj_array = AdjustedArray(
data=label_array,
adjustments={4: [ObjectOverwrite(2, 3, 0, 0, 'ppp')]},
missing_value='',
)
expected_data = array([
['aaa-foo', 'bbb-foo', 'ccc-foo'],
['ddd-foo', 'eee-foo', 'fff-foo'],
['ggg-foo', 'hhh-foo', 'iii-foo'],
['jjj-foo', 'kkk-foo', 'lll-foo'],
['mmm-foo', 'nnn-foo', 'ooo-foo'],
])
expected_label_array = LabelArray(expected_data, missing_value='')
expected_adj_array = AdjustedArray(
data=expected_label_array,
adjustments={4: [ObjectOverwrite(2, 3, 0, 0, 'ppp-foo')]},
missing_value='',
)
adj_array.update_labels(lambda x: x + '-foo')
# Check that the mapped AdjustedArray has the expected baseline
# values and adjustment values.
check_arrays(adj_array.data, expected_adj_array.data)
self.assertEqual(adj_array.adjustments, expected_adj_array.adjustments)
A = Float64Multiply(0, 4, 1, 1, 0.5)
B = Float64Overwrite(3, 3, 4, 4, 4.2)
C = Float64Multiply(0, 2, 0, 0, 0.14)
D = Float64Overwrite(0, 3, 0, 0, 4.0)
E = Float64Overwrite(0, 0, 1, 1, 3.7)
F = Float64Multiply(0, 4, 3, 3, 10.0)
G = Float64Overwrite(5, 5, 4, 4, 1.7)
H = Float64Multiply(0, 4, 2, 2, 0.99)
S = Float64Multiply(0, 1, 4, 4, 5.06)
@parameterized.expand([(
# Initial adjustments
{
1: [A, B],
2: [C],
4: [D],
},
# Adjustments to add
{
1: [E],
2: [F, G],
3: [H, S],
},
# Expected adjustments with 'append'
{
1: [A, B, E],
2: [C, F, G],
3: [H, S],
4: [D],
},
# Expected adjustments with 'prepend'
{
1: [E, A, B],
2: [F, G, C],
3: [H, S],
4: [D],
},
)])
def test_update_adjustments(self, initial_adjustments, adjustments_to_add,
expected_adjustments_with_append,
expected_adjustments_with_prepend):
methods = ['append', 'prepend']
expected_outputs = [
expected_adjustments_with_append, expected_adjustments_with_prepend
]
for method, expected_output in zip(methods, expected_outputs):
data = arange(30, dtype=float).reshape(6, 5)
adjusted_array = AdjustedArray(data, initial_adjustments,
float('nan'))
adjusted_array.update_adjustments(adjustments_to_add, method)
self.assertEqual(adjusted_array.adjustments, expected_output)
def test_adjustments(self):
data = arange(100).reshape(self.ndates, self.nsids)
baseline = DataFrame(data, index=self.dates, columns=self.sids)
# Use the dates from index 10 on and sids 1-3.
dates_slice = slice(10, None, None)
sids_slice = slice(1, 4, None)
# Adjustments that should actually affect the output.
relevant_adjustments = [
{
'sid': 1,
'start_date': None,
'end_date': self.dates[15],
'apply_date': self.dates[16],
'value': 0.5,
'kind': MULTIPLY,
},
{
'sid': 2,
'start_date': self.dates[5],
'end_date': self.dates[15],
'apply_date': self.dates[16],
'value': 1.0,
'kind': ADD,
},
{
'sid': 2,
'start_date': self.dates[15],
'end_date': self.dates[16],
'apply_date': self.dates[17],
'value': 1.0,
'kind': ADD,
},
{
'sid': 3,
'start_date': self.dates[16],
'end_date': self.dates[17],
'apply_date': self.dates[18],
'value': 99.0,
'kind': OVERWRITE,
},
]
# These adjustments shouldn't affect the output.
irrelevant_adjustments = [
{ # Sid Not Requested
'sid': 0,
'start_date': self.dates[16],
'end_date': self.dates[17],
'apply_date': self.dates[18],
'value': -9999.0,
'kind': OVERWRITE,
},
{ # Sid Unknown
'sid': 9999,
'start_date': self.dates[16],
'end_date': self.dates[17],
'apply_date': self.dates[18],
'value': -9999.0,
'kind': OVERWRITE,
},
{ # Date Not Requested
'sid': 2,
'start_date': self.dates[1],
'end_date': self.dates[2],
'apply_date': self.dates[3],
'value': -9999.0,
'kind': OVERWRITE,
},
{ # Date Before Known Data
'sid': 2,
'start_date': self.dates[0] - (2 * trading_day),
'end_date': self.dates[0] - trading_day,
'apply_date': self.dates[0] - trading_day,
'value': -9999.0,
'kind': OVERWRITE,
},
{ # Date After Known Data
'sid': 2,
'start_date': self.dates[-1] + trading_day,
'end_date': self.dates[-1] + (2 * trading_day),
'apply_date': self.dates[-1] + (3 * trading_day),
'value': -9999.0,
'kind': OVERWRITE,
},
]
adjustments = DataFrame(relevant_adjustments + irrelevant_adjustments)
loader = DataFrameLoader(
USEquityPricing.close,
baseline,
adjustments=adjustments,
)
expected_baseline = baseline.iloc[dates_slice, sids_slice]
formatted_adjustments = loader.format_adjustments(
self.dates[dates_slice],
self.sids[sids_slice],
)
expected_formatted_adjustments = {
6: [
Float64Multiply(
first_row=0,
last_row=5,
first_col=0,
last_col=0,
value=0.5,
),
Float64Add(
first_row=0,
last_row=5,
first_col=1,
last_col=1,
value=1.0,
),
],
7: [
Float64Add(
first_row=5,
last_row=6,
first_col=1,
last_col=1,
value=1.0,
),
],
8: [
Float64Overwrite(
first_row=6,
last_row=7,
first_col=2,
last_col=2,
value=99.0,
)
],
}
self.assertEqual(formatted_adjustments, expected_formatted_adjustments)
mask = self.mask[dates_slice, sids_slice]
with patch('zipline.pipeline.loaders.frame.adjusted_array') as m:
loader.load_adjusted_array(
columns=[USEquityPricing.close],
dates=self.dates[dates_slice],
assets=self.sids[sids_slice],
mask=mask,
)
self.assertEqual(m.call_count, 1)
args, kwargs = m.call_args
assert_array_equal(kwargs['data'], expected_baseline.values)
assert_array_equal(kwargs['mask'], mask)
self.assertEqual(kwargs['adjustments'], expected_formatted_adjustments)
def test_ingest(self):
calendar = get_calendar("XNYS")
sessions = calendar.sessions_in_range(self.START_DATE, self.END_DATE)
minutes = calendar.minutes_for_sessions_in_range(
self.START_DATE,
self.END_DATE,
)
sids = tuple(range(3))
equities = make_simple_equity_info(
sids,
self.START_DATE,
self.END_DATE,
)
daily_bar_data = make_bar_data(equities, sessions)
minute_bar_data = make_bar_data(equities, minutes)
first_split_ratio = 0.5
second_split_ratio = 0.1
splits = pd.DataFrame.from_records([
{
"effective_date": str_to_seconds("2014-01-08"),
"ratio": first_split_ratio,
"sid": 0,
},
{
"effective_date": str_to_seconds("2014-01-09"),
"ratio": second_split_ratio,
"sid": 1,
},
])
@self.register(
"bundle",
calendar_name="NYSE",
start_session=self.START_DATE,
end_session=self.END_DATE,
)
def bundle_ingest(
environ,
asset_db_writer,
minute_bar_writer,
daily_bar_writer,
adjustment_writer,
calendar,
start_session,
end_session,
cache,
show_progress,
output_dir,
):
assert environ is self.environ
asset_db_writer.write(equities=equities)
minute_bar_writer.write(minute_bar_data)
daily_bar_writer.write(daily_bar_data)
adjustment_writer.write(splits=splits)
assert isinstance(calendar, TradingCalendar)
assert isinstance(cache, dataframe_cache)
assert isinstance(show_progress, bool)
self.ingest("bundle", environ=self.environ)
bundle = self.load("bundle", environ=self.environ)
assert set(bundle.asset_finder.sids) == set(sids)
columns = "open", "high", "low", "close", "volume"
actual = bundle.equity_minute_bar_reader.load_raw_arrays(
columns,
minutes[0],
minutes[-1],
sids,
)
for actual_column, colname in zip(actual, columns):
np.testing.assert_array_equal(
actual_column,
expected_bar_values_2d(minutes, sids, equities, colname),
err_msg=colname,
)
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
columns,
self.START_DATE,
self.END_DATE,
sids,
)
for actual_column, colname in zip(actual, columns):
np.testing.assert_array_equal(
actual_column,
expected_bar_values_2d(sessions, sids, equities, colname),
err_msg=colname,
)
adjs_for_cols = bundle.adjustment_reader.load_pricing_adjustments(
columns,
sessions,
pd.Index(sids),
)
for column, adjustments in zip(columns, adjs_for_cols[:-1]):
# iterate over all the adjustments but `volume`
assert adjustments == {
2: [
Float64Multiply(
first_row=0,
last_row=2,
first_col=0,
last_col=0,
value=first_split_ratio,
)
],
3: [
Float64Multiply(
first_row=0,
last_row=3,
first_col=1,
last_col=1,
value=second_split_ratio,
)
],
}, column
# check the volume, the value should be 1/ratio
assert adjs_for_cols[-1] == {
2: [
Float64Multiply(
first_row=0,
last_row=2,
first_col=0,
last_col=0,
value=1 / first_split_ratio,
)
],
3: [
Float64Multiply(
first_row=0,
last_row=3,
first_col=1,
last_col=1,
value=1 / second_split_ratio,
)
],
}, "volume"
def _expected_data(self, asset_finder):
sids = {
symbol: asset_finder.lookup_symbol(
symbol,
None,
).sid
for symbol in self.symbols
}
# Load raw data local db.
data = _raw_data(self.symbols, self.start_date, self.end_date,
self.columns)
all_ = data.set_index(
'sid',
append=True,
).unstack()
# fancy list comprehension with statements
@list
@apply
def pricing():
for column in self.columns:
vs = all_[column].values
if column == 'volume':
vs = np.nan_to_num(vs)
yield vs
# the first index our written data will appear in the files on disk
start_idx = (
self.calendar.all_sessions.get_loc(self.start_date, 'ffill') + 1)
######修改到此处
# convert an index into the raw dataframe into an index into the
# final data
i = op.add(start_idx)
def expected_dividend_adjustment(idx, symbol):
sid = sids[symbol]
return (1 -
all_.ix[idx,
('ex_dividend', sid)] / all_.ix[idx - 1,
('close', sid)])
adjustments = [
# ohlc
{
# dividends
i(24): [
Float64Multiply(
first_row=0,
last_row=i(24),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(24, 'AAPL'),
)
],
i(87): [
Float64Multiply(
first_row=0,
last_row=i(87),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(87, 'AAPL'),
)
],
i(150): [
Float64Multiply(
first_row=0,
last_row=i(150),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(150, 'AAPL'),
)
],
i(214): [
Float64Multiply(
first_row=0,
last_row=i(214),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(214, 'AAPL'),
)
],
i(31): [
Float64Multiply(
first_row=0,
last_row=i(31),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(31, 'MSFT'),
)
],
i(90): [
Float64Multiply(
first_row=0,
last_row=i(90),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(90, 'MSFT'),
)
],
i(158): [
Float64Multiply(
first_row=0,
last_row=i(158),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(158, 'MSFT'),
)
],
i(222): [
Float64Multiply(
first_row=0,
last_row=i(222),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(222, 'MSFT'),
)
],
# splits
i(108): [
Float64Multiply(
first_row=0,
last_row=i(108),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=1.0 / 7.0,
)
],
},
] * (len(self.columns) - 1) + [
# volume
{
i(108): [
Float64Multiply(
first_row=0,
last_row=i(108),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=7.0,
)
],
}
]
return pricing, adjustments
def test_ingest(self):
calendar = get_calendar('XNYS')
sessions = calendar.sessions_in_range(self.START_DATE, self.END_DATE)
minutes = calendar.minutes_for_sessions_in_range(
self.START_DATE, self.END_DATE,
)
sids = tuple(range(3))
equities = make_simple_equity_info(
sids,
self.START_DATE,
self.END_DATE,
)
daily_bar_data = make_bar_data(equities, sessions)
minute_bar_data = make_bar_data(equities, minutes)
first_split_ratio = 0.5
second_split_ratio = 0.1
splits = pd.DataFrame.from_records([
{
'effective_date': str_to_seconds('2014-01-08'),
'ratio': first_split_ratio,
'sid': 0,
},
{
'effective_date': str_to_seconds('2014-01-09'),
'ratio': second_split_ratio,
'sid': 1,
},
])
@self.register(
'bundle',
calendar_name='NYSE',
start_session=self.START_DATE,
end_session=self.END_DATE,
)
def bundle_ingest(environ,
asset_db_writer,
minute_bar_writer,
daily_bar_writer,
adjustment_writer,
calendar,
start_session,
end_session,
cache,
output_dir):
assert_is(environ, self.environ)
asset_db_writer.write(equities=equities)
minute_bar_writer.write(minute_bar_data)
daily_bar_writer.write(daily_bar_data)
adjustment_writer.write(splits=splits)
assert_is_instance(calendar, TradingCalendar)
assert_is_instance(cache, dataframe_cache)
self.ingest('bundle', environ=self.environ)
bundle = self.load('bundle', environ=self.environ)
assert_equal(set(bundle.asset_finder.sids), set(sids))
columns = 'open', 'high', 'low', 'close', 'volume'
actual = bundle.equity_minute_bar_reader.load_raw_arrays(
columns,
minutes[0],
minutes[-1],
sids,
)
for actual_column, colname in zip(actual, columns):
assert_equal(
actual_column,
expected_bar_values_2d(minutes, sids, equities, colname),
msg=colname,
)
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
columns,
self.START_DATE,
self.END_DATE,
sids,
)
for actual_column, colname in zip(actual, columns):
assert_equal(
actual_column,
expected_bar_values_2d(sessions, sids, equities, colname),
msg=colname,
)
adjs_for_cols = bundle.adjustment_reader.load_pricing_adjustments(
columns,
sessions,
pd.Index(sids),
)
for column, adjustments in zip(columns, adjs_for_cols[:-1]):
# iterate over all the adjustments but `volume`
assert_equal(
adjustments,
{
2: [Float64Multiply(
first_row=0,
last_row=2,
first_col=0,
last_col=0,
value=first_split_ratio,
)],
3: [Float64Multiply(
first_row=0,
last_row=3,
first_col=1,
last_col=1,
value=second_split_ratio,
)],
},
msg=column,
)
# check the volume, the value should be 1/ratio
assert_equal(
adjs_for_cols[-1],
{
2: [Float64Multiply(
first_row=0,
last_row=2,
first_col=0,
last_col=0,
value=1 / first_split_ratio,
)],
3: [Float64Multiply(
first_row=0,
last_row=3,
first_col=1,
last_col=1,
value=1 / second_split_ratio,
)],
},
msg='volume',
)
def _expected_data(self, asset_finder):
sids = {
symbol: asset_finder.lookup_symbol(
symbol,
self.asset_start,
).sid
for symbol in self.symbols
}
def per_symbol(symbol):
df = pd.read_csv(
test_resource_path('quandl_samples', symbol + '.csv.gz'),
parse_dates=['Date'],
index_col='Date',
usecols=[
'Open',
'High',
'Low',
'Close',
'Volume',
'Date',
'Ex-Dividend',
'Split Ratio',
],
na_values=['NA'],
).rename(columns={
'Open': 'open',
'High': 'high',
'Low': 'low',
'Close': 'close',
'Volume': 'volume',
'Date': 'date',
'Ex-Dividend': 'ex_dividend',
'Split Ratio': 'split_ratio',
})
df['sid'] = sids[symbol]
return df
all_ = pd.concat(map(per_symbol, self.symbols)).set_index(
'sid',
append=True,
).unstack()
# fancy list comprehension with statements
@list
@apply
def pricing():
for column in self.columns:
vs = all_[column].values
if column == 'volume':
vs = np.nan_to_num(vs)
yield vs
# the first index our written data will appear in the files on disk
start_idx = self.calendar.get_loc(self.asset_start, 'ffill') + 1
# convert an index into the raw dataframe into an index into the
# final data
i = op.add(start_idx)
def expected_dividend_adjustment(idx, symbol):
sid = sids[symbol]
return (
1 -
all_.ix[idx, ('ex_dividend', sid)] /
all_.ix[idx - 1, ('close', sid)]
)
adjustments = [
# ohlc
{
# dividends
i(24): [Float64Multiply(
first_row=0,
last_row=i(24),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(24, 'AAPL'),
)],
i(87): [Float64Multiply(
first_row=0,
last_row=i(87),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(87, 'AAPL'),
)],
i(150): [Float64Multiply(
first_row=0,
last_row=i(150),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(150, 'AAPL'),
)],
i(214): [Float64Multiply(
first_row=0,
last_row=i(214),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=expected_dividend_adjustment(214, 'AAPL'),
)],
i(31): [Float64Multiply(
first_row=0,
last_row=i(31),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(31, 'MSFT'),
)],
i(90): [Float64Multiply(
first_row=0,
last_row=i(90),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(90, 'MSFT'),
)],
i(222): [Float64Multiply(
first_row=0,
last_row=i(222),
first_col=sids['MSFT'],
last_col=sids['MSFT'],
value=expected_dividend_adjustment(222, 'MSFT'),
)],
# splits
i(108): [Float64Multiply(
first_row=0,
last_row=i(108),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=1.0 / 7.0,
)],
},
] * (len(self.columns) - 1) + [
# volume
{
i(108): [Float64Multiply(
first_row=0,
last_row=i(108),
first_col=sids['AAPL'],
last_col=sids['AAPL'],
value=7.0,
)],
}
]
return pricing, adjustments
def _expected_data(self, asset_finder):
sids = {
symbol: asset_finder.lookup_symbol(
symbol,
None,
).sid
for symbol in self.symbols
}
# Load raw data from quandl test resources.
data = load_data_table(
file=test_resource_path("quandl_samples", "QUANDL_ARCHIVE.zip"),
index_col="date",
)
data["sid"] = pd.factorize(data.symbol)[0]
all_ = data.set_index(
"sid",
append=True,
).unstack()
# fancy list comprehension with statements
@list
@apply
def pricing():
for column in self.columns:
vs = all_[column].values
if column == "volume":
vs = np.nan_to_num(vs)
yield vs
# the first index our written data will appear in the files on disk
start_idx = self.calendar.all_sessions.get_loc(self.start_date,
"ffill") + 1
# convert an index into the raw dataframe into an index into the
# final data
i = op.add(start_idx)
def expected_dividend_adjustment(idx, symbol):
sid = sids[symbol]
return (1 - all_.iloc[idx]["ex_dividend", sid] /
all_.iloc[idx - 1]["close", sid])
adjustments = [
# ohlc
{
# dividends
i(24): [
Float64Multiply(
first_row=0,
last_row=i(24),
first_col=sids["AAPL"],
last_col=sids["AAPL"],
value=expected_dividend_adjustment(24, "AAPL"),
)
],
i(87): [
Float64Multiply(
first_row=0,
last_row=i(87),
first_col=sids["AAPL"],
last_col=sids["AAPL"],
value=expected_dividend_adjustment(87, "AAPL"),
)
],
i(150): [
Float64Multiply(
first_row=0,
last_row=i(150),
first_col=sids["AAPL"],
last_col=sids["AAPL"],
value=expected_dividend_adjustment(150, "AAPL"),
)
],
i(214): [
Float64Multiply(
first_row=0,
last_row=i(214),
first_col=sids["AAPL"],
last_col=sids["AAPL"],
value=expected_dividend_adjustment(214, "AAPL"),
)
],
i(31): [
Float64Multiply(
first_row=0,
last_row=i(31),
first_col=sids["MSFT"],
last_col=sids["MSFT"],
value=expected_dividend_adjustment(31, "MSFT"),
)
],
i(90): [
Float64Multiply(
first_row=0,
last_row=i(90),
first_col=sids["MSFT"],
last_col=sids["MSFT"],
value=expected_dividend_adjustment(90, "MSFT"),
)
],
i(158): [
Float64Multiply(
first_row=0,
last_row=i(158),
first_col=sids["MSFT"],
last_col=sids["MSFT"],
value=expected_dividend_adjustment(158, "MSFT"),
)
],
i(222): [
Float64Multiply(
first_row=0,
last_row=i(222),
first_col=sids["MSFT"],
last_col=sids["MSFT"],
value=expected_dividend_adjustment(222, "MSFT"),
)
],
# splits
i(108): [
Float64Multiply(
first_row=0,
last_row=i(108),
first_col=sids["AAPL"],
last_col=sids["AAPL"],
value=1.0 / 7.0,
)
],
},
] * (len(self.columns) - 1) + [
# volume
{
i(108): [
Float64Multiply(
first_row=0,
last_row=i(108),
first_col=sids["AAPL"],
last_col=sids["AAPL"],
value=7.0,
)
],
}
]
return pricing, adjustments
def _get_adjustments_in_range(self, asset, dts, field,
is_perspective_after):
"""
Get the Float64Multiply objects to pass to an AdjustedArrayWindow.
For the use of AdjustedArrayWindow in the loader, which looks back
from current simulation time back to a window of data the dictionary is
structured with:
- the key into the dictionary for adjustments is the location of the
day from which the window is being viewed.
- the start of all multiply objects is always 0 (in each window all
adjustments are overlapping)
- the end of the multiply object is the location before the calendar
location of the adjustment action, making all days before the event
adjusted.
Parameters
----------
asset : Asset
The assets for which to get adjustments.
days : iterable of datetime64-like
The days for which adjustment data is needed.
field : str
OHLCV field for which to get the adjustments.
is_perspective_after : bool
see: `USEquityHistoryLoader.history`
If True, the index at which the Multiply object is registered to
be popped is calculated so that it applies to the last slot in the
sliding window when the adjustment occurs immediately after the dt
that slot represents.
Returns
-------
out : The adjustments as a dict of loc -> Float64Multiply
"""
sid = int(asset)
start = normalize_date(dts[0])
end = normalize_date(dts[-1])
adjs = {}
if field != 'volume':
mergers = self._adjustments_reader.get_adjustments_for_sid(
'mergers', sid)
for m in mergers:
dt = m[0]
if start < dt <= end:
end_loc = dts.searchsorted(dt)
adj_loc = end_loc
if is_perspective_after:
# Set adjustment pop location so that it applies
# to last value if adjustment occurs immediately after
# the last slot.
adj_loc -= 1
mult = Float64Multiply(0,
end_loc - 1,
0,
0,
m[1])
try:
adjs[adj_loc].append(mult)
except KeyError:
adjs[adj_loc] = [mult]
divs = self._adjustments_reader.get_adjustments_for_sid(
'dividends', sid)
for d in divs:
dt = d[0]
if start < dt <= end:
end_loc = dts.searchsorted(dt)
adj_loc = end_loc
if is_perspective_after:
# Set adjustment pop location so that it applies
# to last value if adjustment occurs immediately after
# the last slot.
adj_loc -= 1
mult = Float64Multiply(0,
end_loc - 1,
0,
0,
d[1])
try:
adjs[adj_loc].append(mult)
except KeyError:
adjs[adj_loc] = [mult]
splits = self._adjustments_reader.get_adjustments_for_sid(
'splits', sid)
for s in splits:
dt = s[0]
if start < dt <= end:
if field == 'volume':
ratio = 1.0 / s[1]
else:
ratio = s[1]
end_loc = dts.searchsorted(dt)
adj_loc = end_loc
if is_perspective_after:
# Set adjustment pop location so that it applies
# to last value if adjustment occurs immediately after
# the last slot.
adj_loc -= 1
mult = Float64Multiply(0,
end_loc - 1,
0,
0,
ratio)
try:
adjs[adj_loc].append(mult)
except KeyError:
adjs[adj_loc] = [mult]
return adjs
def test_ingest(self):
start = pd.Timestamp('2014-01-06', tz='utc')
end = pd.Timestamp('2014-01-10', tz='utc')
trading_days = get_calendar('NYSE').all_trading_days
calendar = trading_days[trading_days.slice_indexer(start, end)]
minutes = get_calendar('NYSE').trading_minutes_for_days_in_range(
calendar[0], calendar[-1])
sids = tuple(range(3))
equities = make_simple_equity_info(
sids,
calendar[0],
calendar[-1],
)
daily_bar_data = make_bar_data(equities, calendar)
minute_bar_data = make_bar_data(equities, minutes)
first_split_ratio = 0.5
second_split_ratio = 0.1
splits = pd.DataFrame.from_records([
{
'effective_date': str_to_seconds('2014-01-08'),
'ratio': first_split_ratio,
'sid': 0,
},
{
'effective_date': str_to_seconds('2014-01-09'),
'ratio': second_split_ratio,
'sid': 1,
},
])
schedule = get_calendar('NYSE').schedule
@self.register(
'bundle',
calendar=calendar,
opens=schedule.market_open[calendar[0]:calendar[-1]],
closes=schedule.market_close[calendar[0]:calendar[-1]],
)
def bundle_ingest(environ, asset_db_writer, minute_bar_writer,
daily_bar_writer, adjustment_writer, calendar, cache,
show_progress, output_dir):
assert_is(environ, self.environ)
asset_db_writer.write(equities=equities)
minute_bar_writer.write(minute_bar_data)
daily_bar_writer.write(daily_bar_data)
adjustment_writer.write(splits=splits)
assert_is_instance(calendar, pd.DatetimeIndex)
assert_is_instance(cache, dataframe_cache)
assert_is_instance(show_progress, bool)
self.ingest('bundle', environ=self.environ)
bundle = self.load('bundle', environ=self.environ)
assert_equal(set(bundle.asset_finder.sids), set(sids))
columns = 'open', 'high', 'low', 'close', 'volume'
actual = bundle.equity_minute_bar_reader.load_raw_arrays(
columns,
minutes[0],
minutes[-1],
sids,
)
for actual_column, colname in zip(actual, columns):
assert_equal(
actual_column,
expected_bar_values_2d(minutes, equities, colname),
msg=colname,
)
actual = bundle.equity_daily_bar_reader.load_raw_arrays(
columns,
calendar[0],
calendar[-1],
sids,
)
for actual_column, colname in zip(actual, columns):
assert_equal(
actual_column,
expected_bar_values_2d(calendar, equities, colname),
msg=colname,
)
adjustments_for_cols = bundle.adjustment_reader.load_adjustments(
columns,
calendar,
pd.Index(sids),
)
for column, adjustments in zip(columns, adjustments_for_cols[:-1]):
# iterate over all the adjustments but `volume`
assert_equal(
adjustments,
{
2: [
Float64Multiply(
first_row=0,
last_row=2,
first_col=0,
last_col=0,
value=first_split_ratio,
)
],
3: [
Float64Multiply(
first_row=0,
last_row=3,
first_col=1,
last_col=1,
value=second_split_ratio,
)
],
},
msg=column,
)
# check the volume, the value should be 1/ratio
assert_equal(
adjustments_for_cols[-1],
{
2: [
Float64Multiply(
first_row=0,
last_row=2,
first_col=0,
last_col=0,
value=1 / first_split_ratio,
)
],
3: [
Float64Multiply(
first_row=0,
last_row=3,
first_col=1,
last_col=1,
value=1 / second_split_ratio,
)
],
},
msg='volume',
)
class TestAdjustedArray:
def test_traverse_invalidating(self):
data = np.arange(5 * 3, dtype="f8").reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float("nan"))
for _ in adjusted_array.traverse(1, copy=False):
pass
assert_equal(data, original_data * 2)
err_msg = "cannot traverse invalidated AdjustedArray"
with pytest.raises(ValueError, match=err_msg):
adjusted_array.traverse(1)
def test_copy(self):
data = np.arange(5 * 3, dtype="f8").reshape(5, 3)
original_data = data.copy()
adjustments = {2: [Float64Multiply(0, 4, 0, 2, 2.0)]}
adjusted_array = AdjustedArray(data, adjustments, float("nan"))
traverse_copy = adjusted_array.copy()
clean_copy = adjusted_array.copy()
a_it = adjusted_array.traverse(2, copy=False)
b_it = traverse_copy.traverse(2, copy=False)
for a, b in zip(a_it, b_it):
assert_equal(a, b)
err_msg = "cannot copy invalidated AdjustedArray"
with pytest.raises(ValueError, match=err_msg):
adjusted_array.copy()
# the clean copy should have the original data even though the
# original adjusted array has it's data mutated in place
assert_equal(clean_copy.data, original_data)
assert_equal(adjusted_array.data, original_data * 2)
@pytest.mark.parametrize(
"""name, data, lookback, adjustments, missing_value,\
perspective_offset, expected_output""",
chain(
_gen_unadjusted_cases(
"float",
make_input=as_dtype(float64_dtype),
make_expected_output=as_dtype(float64_dtype),
missing_value=default_missing_value_for_dtype(float64_dtype),
),
_gen_unadjusted_cases(
"datetime",
make_input=as_dtype(datetime64ns_dtype),
make_expected_output=as_dtype(datetime64ns_dtype),
missing_value=default_missing_value_for_dtype(datetime64ns_dtype),
),
# Test passing an array of strings to AdjustedArray.
_gen_unadjusted_cases(
"bytes_ndarray",
make_input=as_dtype(bytes_dtype),
make_expected_output=as_labelarray(bytes_dtype, b""),
missing_value=b"",
),
_gen_unadjusted_cases(
"unicode_ndarray",
make_input=as_dtype(unicode_dtype),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value="",
),
_gen_unadjusted_cases(
"object_ndarray",
make_input=lambda a: a.astype(unicode).astype(object),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value="",
),
# Test passing a LabelArray directly to AdjustedArray.
_gen_unadjusted_cases(
"bytes_labelarray",
make_input=as_labelarray(bytes_dtype, b""),
make_expected_output=as_labelarray(bytes_dtype, b""),
missing_value=b"",
),
_gen_unadjusted_cases(
"unicode_labelarray",
make_input=as_labelarray(unicode_dtype, None),
make_expected_output=as_labelarray(unicode_dtype, None),
missing_value="",
),
_gen_unadjusted_cases(
"object_labelarray",
make_input=(lambda a: LabelArray(a.astype(unicode).astype(object), "")),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value="",
),
),
)
def test_no_adjustments(
self,
name,
data,
lookback,
adjustments,
missing_value,
perspective_offset,
expected_output,
):
array = AdjustedArray(data, adjustments, missing_value)
for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable.
in_out = zip(array.traverse(lookback), expected_output)
for yielded, expected_yield in in_out:
check_arrays(yielded, expected_yield)
@pytest.mark.parametrize(
"name, data, lookback, adjustments, missing_value,\
perspective_offset, expected",
_gen_multiplicative_adjustment_cases(float64_dtype),
)
def test_multiplicative_adjustments(
self,
name,
data,
lookback,
adjustments,
missing_value,
perspective_offset,
expected,
):
array = AdjustedArray(data, adjustments, missing_value)
for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable.
window_iter = array.traverse(
lookback,
perspective_offset=perspective_offset,
)
for yielded, expected_yield in zip_longest(window_iter, expected):
check_arrays(yielded, expected_yield)
@pytest.mark.parametrize(
"name, baseline, lookback, adjustments,\
missing_value, perspective_offset, expected",
chain(
_gen_overwrite_adjustment_cases(bool_dtype),
_gen_overwrite_adjustment_cases(int64_dtype),
_gen_overwrite_adjustment_cases(float64_dtype),
_gen_overwrite_adjustment_cases(datetime64ns_dtype),
_gen_overwrite_1d_array_adjustment_case(float64_dtype),
_gen_overwrite_1d_array_adjustment_case(datetime64ns_dtype),
_gen_overwrite_1d_array_adjustment_case(bool_dtype),
# There are six cases here:
# Using np.bytes/np.unicode/object arrays as inputs.
# Passing np.bytes/np.unicode/object arrays to LabelArray,
# and using those as input.
#
# The outputs should always be LabelArrays.
_gen_unadjusted_cases(
"bytes_ndarray",
make_input=as_dtype(bytes_dtype),
make_expected_output=as_labelarray(bytes_dtype, b""),
missing_value=b"",
),
_gen_unadjusted_cases(
"unicode_ndarray",
make_input=as_dtype(unicode_dtype),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value="",
),
_gen_unadjusted_cases(
"object_ndarray",
make_input=lambda a: a.astype(unicode).astype(object),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value="",
),
_gen_unadjusted_cases(
"bytes_labelarray",
make_input=as_labelarray(bytes_dtype, b""),
make_expected_output=as_labelarray(bytes_dtype, b""),
missing_value=b"",
),
_gen_unadjusted_cases(
"unicode_labelarray",
make_input=as_labelarray(unicode_dtype, ""),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value="",
),
_gen_unadjusted_cases(
"object_labelarray",
make_input=(
lambda a: LabelArray(
a.astype(unicode).astype(object),
None,
)
),
make_expected_output=as_labelarray(unicode_dtype, ""),
missing_value=None,
),
),
)
def test_overwrite_adjustment_cases(
self,
name,
baseline,
lookback,
adjustments,
missing_value,
perspective_offset,
expected,
):
array = AdjustedArray(baseline, adjustments, missing_value)
for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable.
window_iter = array.traverse(
lookback,
perspective_offset=perspective_offset,
)
for yielded, expected_yield in zip_longest(window_iter, expected):
check_arrays(yielded, expected_yield)
def test_object1darrayoverwrite(self):
pairs = [u + l for u, l in product(ascii_uppercase, ascii_lowercase)]
categories = pairs + ["~" + c for c in pairs]
baseline = LabelArray(
np.array([["".join((r, c)) for c in "abc"] for r in ascii_uppercase]),
None,
categories,
)
full_expected = baseline.copy()
def flip(cs):
if cs is None:
return None
if cs[0] != "~":
return "~" + cs
return cs
def make_overwrite(fr, lr, fc, lc):
fr, lr, fc, lc = map(ord, (fr, lr, fc, lc))
fr -= ord("A")
lr -= ord("A")
fc -= ord("a")
lc -= ord("a")
return Object1DArrayOverwrite(
fr,
lr,
fc,
lc,
baseline[fr : lr + 1, fc].map(flip),
)
overwrites = {
3: [make_overwrite("A", "B", "a", "a")],
4: [make_overwrite("A", "C", "b", "c")],
5: [make_overwrite("D", "D", "a", "b")],
}
it = AdjustedArray(baseline, overwrites, None).traverse(3)
window = next(it)
expected = full_expected[:3]
check_arrays(window, expected)
window = next(it)
full_expected[0:2, 0] = LabelArray(["~Aa", "~Ba"], None)
expected = full_expected[1:4]
check_arrays(window, expected)
window = next(it)
full_expected[0:3, 1:3] = LabelArray(
[["~Ab", "~Ac"], ["~Bb", "~Bc"], ["~Cb", "~Cb"]], None
)
expected = full_expected[2:5]
check_arrays(window, expected)
window = next(it)
full_expected[3, :2] = "~Da"
expected = full_expected[3:6]
check_arrays(window, expected)
def test_invalid_lookback(self):
data = np.arange(30, dtype=float).reshape(6, 5)
adj_array = AdjustedArray(data, {}, float("nan"))
with pytest.raises(WindowLengthTooLong):
adj_array.traverse(7)
with pytest.raises(WindowLengthNotPositive):
adj_array.traverse(0)
with pytest.raises(WindowLengthNotPositive):
adj_array.traverse(-1)
def test_array_views_arent_writable(self):
data = np.arange(30, dtype=float).reshape(6, 5)
adj_array = AdjustedArray(data, {}, float("nan"))
for frame in adj_array.traverse(3):
with pytest.raises(ValueError):
frame[0, 0] = 5.0
def test_inspect(self):
data = np.arange(15, dtype=float).reshape(5, 3)
adj_array = AdjustedArray(
data,
{4: [Float64Multiply(2, 3, 0, 0, 4.0)]},
float("nan"),
)
# TODO: CHECK WHY DO I NEED TO FIX THE INDENT IN THE EXPECTED?
expected = dedent(
"""\
Adjusted Array (float64):
Data:
array([[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.],
[ 9., 10., 11.],
[12., 13., 14.]])
Adjustments:
{4: [Float64Multiply(first_row=2, last_row=3, first_col=0, \
last_col=0, value=4.000000)]}
"""
)
got = adj_array.inspect()
assert expected == got
def test_update_labels(self):
data = np.array(
[
["aaa", "bbb", "ccc"],
["ddd", "eee", "fff"],
["ggg", "hhh", "iii"],
["jjj", "kkk", "lll"],
["mmm", "nnn", "ooo"],
]
)
label_array = LabelArray(data, missing_value="")
adj_array = AdjustedArray(
data=label_array,
adjustments={4: [ObjectOverwrite(2, 3, 0, 0, "ppp")]},
missing_value="",
)
expected_data = np.array(
[
["aaa-foo", "bbb-foo", "ccc-foo"],
["ddd-foo", "eee-foo", "fff-foo"],
["ggg-foo", "hhh-foo", "iii-foo"],
["jjj-foo", "kkk-foo", "lll-foo"],
["mmm-foo", "nnn-foo", "ooo-foo"],
]
)
expected_label_array = LabelArray(expected_data, missing_value="")
expected_adj_array = AdjustedArray(
data=expected_label_array,
adjustments={4: [ObjectOverwrite(2, 3, 0, 0, "ppp-foo")]},
missing_value="",
)
adj_array.update_labels(lambda x: x + "-foo")
# Check that the mapped AdjustedArray has the expected baseline
# values and adjustment values.
check_arrays(adj_array.data, expected_adj_array.data)
assert adj_array.adjustments == expected_adj_array.adjustments
A = Float64Multiply(0, 4, 1, 1, 0.5)
B = Float64Overwrite(3, 3, 4, 4, 4.2)
C = Float64Multiply(0, 2, 0, 0, 0.14)
D = Float64Overwrite(0, 3, 0, 0, 4.0)
E = Float64Overwrite(0, 0, 1, 1, 3.7)
F = Float64Multiply(0, 4, 3, 3, 10.0)
G = Float64Overwrite(5, 5, 4, 4, 1.7)
H = Float64Multiply(0, 4, 2, 2, 0.99)
S = Float64Multiply(0, 1, 4, 4, 5.06)
@pytest.mark.parametrize(
"initial_adjustments, adjustments_to_add,\
expected_adjustments_with_append, expected_adjustments_with_prepend",
[
(
# Initial adjustments
{
1: [A, B],
2: [C],
4: [D],
},
# Adjustments to add
{
1: [E],
2: [F, G],
3: [H, S],
},
# Expected adjustments with 'append'
{
1: [A, B, E],
2: [C, F, G],
3: [H, S],
4: [D],
},
# Expected adjustments with 'prepend'
{
1: [E, A, B],
2: [F, G, C],
3: [H, S],
4: [D],
},
)
],
)
def test_update_adjustments(
self,
initial_adjustments,
adjustments_to_add,
expected_adjustments_with_append,
expected_adjustments_with_prepend,
):
methods = ["append", "prepend"]
expected_outputs = [
expected_adjustments_with_append,
expected_adjustments_with_prepend,
]
for method, expected_output in zip(methods, expected_outputs):
data = np.arange(30, dtype=float).reshape(6, 5)
adjusted_array = AdjustedArray(data, initial_adjustments, float("nan"))
adjusted_array.update_adjustments(adjustments_to_add, method)
assert adjusted_array.adjustments == expected_output
def test_adjustments(self):
data = np.arange(100).reshape(self.ndates, self.nsids)
baseline = pd.DataFrame(data, index=self.dates, columns=self.sids)
# Use the dates from index 10 on and sids 1-3.
dates_slice = slice(10, None, None)
sids_slice = slice(1, 4, None)
# Adjustments that should actually affect the output.
relevant_adjustments = [
{
"sid": 1,
"start_date": None,
"end_date": self.dates[15],
"apply_date": self.dates[16],
"value": 0.5,
"kind": MULTIPLY,
},
{
"sid": 2,
"start_date": self.dates[5],
"end_date": self.dates[15],
"apply_date": self.dates[16],
"value": 1.0,
"kind": ADD,
},
{
"sid": 2,
"start_date": self.dates[15],
"end_date": self.dates[16],
"apply_date": self.dates[17],
"value": 1.0,
"kind": ADD,
},
{
"sid": 3,
"start_date": self.dates[16],
"end_date": self.dates[17],
"apply_date": self.dates[18],
"value": 99.0,
"kind": OVERWRITE,
},
]
# These adjustments shouldn't affect the output.
irrelevant_adjustments = [
{ # Sid Not Requested
"sid": 0,
"start_date": self.dates[16],
"end_date": self.dates[17],
"apply_date": self.dates[18],
"value": -9999.0,
"kind": OVERWRITE,
},
{ # Sid Unknown
"sid": 9999,
"start_date": self.dates[16],
"end_date": self.dates[17],
"apply_date": self.dates[18],
"value": -9999.0,
"kind": OVERWRITE,
},
{ # Date Not Requested
"sid": 2,
"start_date": self.dates[1],
"end_date": self.dates[2],
"apply_date": self.dates[3],
"value": -9999.0,
"kind": OVERWRITE,
},
{ # Date Before Known Data
"sid": 2,
"start_date": self.dates[0] - (2 * self.trading_day),
"end_date": self.dates[0] - self.trading_day,
"apply_date": self.dates[0] - self.trading_day,
"value": -9999.0,
"kind": OVERWRITE,
},
{ # Date After Known Data
"sid": 2,
"start_date": self.dates[-1] + self.trading_day,
"end_date": self.dates[-1] + (2 * self.trading_day),
"apply_date": self.dates[-1] + (3 * self.trading_day),
"value": -9999.0,
"kind": OVERWRITE,
},
]
adjustments = pd.DataFrame(relevant_adjustments +
irrelevant_adjustments)
loader = DataFrameLoader(
USEquityPricing.close,
baseline,
adjustments=adjustments,
)
expected_baseline = baseline.iloc[dates_slice, sids_slice]
formatted_adjustments = loader.format_adjustments(
self.dates[dates_slice],
self.sids[sids_slice],
)
expected_formatted_adjustments = {
6: [
Float64Multiply(
first_row=0,
last_row=5,
first_col=0,
last_col=0,
value=0.5,
),
Float64Add(
first_row=0,
last_row=5,
first_col=1,
last_col=1,
value=1.0,
),
],
7: [
Float64Add(
first_row=5,
last_row=6,
first_col=1,
last_col=1,
value=1.0,
),
],
8: [
Float64Overwrite(
first_row=6,
last_row=7,
first_col=2,
last_col=2,
value=99.0,
)
],
}
assert formatted_adjustments == expected_formatted_adjustments
mask = self.mask[dates_slice, sids_slice]
with patch("zipline.pipeline.loaders.frame.AdjustedArray") as m:
loader.load_adjusted_array(
US_EQUITIES,
columns=[USEquityPricing.close],
dates=self.dates[dates_slice],
sids=self.sids[sids_slice],
mask=mask,
)
assert m.call_count == 1
args, kwargs = m.call_args
assert_array_equal(kwargs["data"], expected_baseline.values)
assert kwargs["adjustments"] == expected_formatted_adjustments
def _get_adjustments_in_range(self, asset, dts, field):
"""
Get the Float64Multiply objects to pass to an AdjustedArrayWindow.
For the use of AdjustedArrayWindow in the loader, which looks back
from current simulation time back to a window of data the dictionary is
structured with:
- the key into the dictionary for adjustments is the location of the
day from which the window is being viewed.
- the start of all multiply objects is always 0 (in each window all
adjustments are overlapping)
- the end of the multiply object is the location before the calendar
location of the adjustment action, making all days before the event
adjusted.
Parameters
----------
asset : Asset
The assets for which to get adjustments.
dts : iterable of datetime64-like
The dts for which adjustment data is needed.
field : str
OHLCV field for which to get the adjustments.
Returns
-------
out : dict[loc -> Float64Multiply]
The adjustments as a dict of loc -> Float64Multiply
"""
sid = int(asset)
start = dts[0].normalize()
end = dts[-1].normalize()
adjs = {}
if field != 'volume':
mergers = self._adjustments_reader.get_adjustments_for_sid(
'mergers', sid)
for m in mergers:
dt = m[0]
if start < dt <= end:
end_loc = dts.searchsorted(dt)
adj_loc = end_loc
mult = Float64Multiply(0, end_loc - 1, 0, 0, m[1])
try:
adjs[adj_loc].append(mult)
except KeyError:
adjs[adj_loc] = [mult]
divs = self._adjustments_reader.get_adjustments_for_sid(
'dividends', sid)
for d in divs:
dt = d[0]
if start < dt <= end:
end_loc = dts.searchsorted(dt)
adj_loc = end_loc
mult = Float64Multiply(0, end_loc - 1, 0, 0, d[1])
try:
adjs[adj_loc].append(mult)
except KeyError:
adjs[adj_loc] = [mult]
splits = self._adjustments_reader.get_adjustments_for_sid(
'splits', sid)
for s in splits:
dt = s[0]
if start < dt <= end:
if field == 'volume':
ratio = 1.0 / s[1]
else:
ratio = s[1]
end_loc = dts.searchsorted(dt)
adj_loc = end_loc
mult = Float64Multiply(0, end_loc - 1, 0, 0, ratio)
try:
adjs[adj_loc].append(mult)
except KeyError:
adjs[adj_loc] = [mult]
return adjs
