mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-13 12:00:16 +08:00
TST: tests for new forward fill
This commit is contained in:
+251
-69
@@ -295,8 +295,8 @@ class BlazeToPipelineTestCase(TestCase):
|
|||||||
loader=self.garbage_loader,
|
loader=self.garbage_loader,
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_id(self):
|
def _test_id(self, df, dshape, expected, finder, add):
|
||||||
expr = bz.Data(self.df, name='expr', dshape=self.dshape)
|
expr = bz.Data(df, name='expr', dshape=dshape)
|
||||||
loader = BlazeLoader()
|
loader = BlazeLoader()
|
||||||
ds = from_blaze(
|
ds = from_blaze(
|
||||||
expr,
|
expr,
|
||||||
@@ -304,7 +304,8 @@ class BlazeToPipelineTestCase(TestCase):
|
|||||||
no_deltas_rule=no_deltas_rules.ignore,
|
no_deltas_rule=no_deltas_rules.ignore,
|
||||||
)
|
)
|
||||||
p = Pipeline()
|
p = Pipeline()
|
||||||
p.add(ds.value.latest, 'value')
|
for a in add:
|
||||||
|
p.add(getattr(ds, a).latest, a)
|
||||||
dates = self.dates
|
dates = self.dates
|
||||||
|
|
||||||
with tmp_asset_finder() as finder:
|
with tmp_asset_finder() as finder:
|
||||||
@@ -314,13 +315,6 @@ class BlazeToPipelineTestCase(TestCase):
|
|||||||
finder,
|
finder,
|
||||||
).run_pipeline(p, dates[0], dates[-1])
|
).run_pipeline(p, dates[0], dates[-1])
|
||||||
|
|
||||||
expected = self.df.drop('asof_date', axis=1).set_index(
|
|
||||||
['timestamp', 'sid'],
|
|
||||||
)
|
|
||||||
expected.index = pd.MultiIndex.from_product((
|
|
||||||
expected.index.levels[0],
|
|
||||||
finder.retrieve_all(expected.index.levels[1]),
|
|
||||||
))
|
|
||||||
assert_frame_equal(result, expected, check_dtype=False)
|
assert_frame_equal(result, expected, check_dtype=False)
|
||||||
|
|
||||||
def test_custom_query_time_tz(self):
|
def test_custom_query_time_tz(self):
|
||||||
@@ -360,75 +354,263 @@ class BlazeToPipelineTestCase(TestCase):
|
|||||||
))
|
))
|
||||||
assert_frame_equal(result, expected, check_dtype=False)
|
assert_frame_equal(result, expected, check_dtype=False)
|
||||||
|
|
||||||
|
def test_id(self):
|
||||||
|
"""
|
||||||
|
input (self.df):
|
||||||
|
asof_date sid timestamp value
|
||||||
|
0 2014-01-01 65 2014-01-01 0
|
||||||
|
1 2014-01-01 66 2014-01-01 1
|
||||||
|
2 2014-01-01 67 2014-01-01 2
|
||||||
|
3 2014-01-02 65 2014-01-02 1
|
||||||
|
4 2014-01-02 66 2014-01-02 2
|
||||||
|
5 2014-01-02 67 2014-01-02 3
|
||||||
|
6 2014-01-03 65 2014-01-03 2
|
||||||
|
7 2014-01-03 66 2014-01-03 3
|
||||||
|
8 2014-01-03 67 2014-01-03 4
|
||||||
|
|
||||||
|
output (expected)
|
||||||
|
value
|
||||||
|
2014-01-01 Equity(65 [A]) 0
|
||||||
|
Equity(66 [B]) 1
|
||||||
|
Equity(67 [C]) 2
|
||||||
|
2014-01-02 Equity(65 [A]) 1
|
||||||
|
Equity(66 [B]) 2
|
||||||
|
Equity(67 [C]) 3
|
||||||
|
2014-01-03 Equity(65 [A]) 2
|
||||||
|
Equity(66 [B]) 3
|
||||||
|
Equity(67 [C]) 4
|
||||||
|
"""
|
||||||
|
with tmp_asset_finder() as finder:
|
||||||
|
expected = self.df.drop('asof_date', axis=1).set_index(
|
||||||
|
['timestamp', 'sid'],
|
||||||
|
)
|
||||||
|
expected.index = pd.MultiIndex.from_product((
|
||||||
|
expected.index.levels[0],
|
||||||
|
finder.retrieve_all(expected.index.levels[1]),
|
||||||
|
))
|
||||||
|
self._test_id(self.df, self.dshape, expected, finder, ('value',))
|
||||||
|
|
||||||
|
def test_id_ffill_out_of_window(self):
|
||||||
|
"""
|
||||||
|
input (df):
|
||||||
|
|
||||||
|
asof_date timestamp sid other value
|
||||||
|
0 2013-12-22 2013-12-22 65 0 0
|
||||||
|
1 2013-12-22 2013-12-22 66 NaN 1
|
||||||
|
2 2013-12-22 2013-12-22 67 2 NaN
|
||||||
|
3 2013-12-23 2013-12-23 65 NaN 1
|
||||||
|
4 2013-12-23 2013-12-23 66 2 NaN
|
||||||
|
5 2013-12-23 2013-12-23 67 3 3
|
||||||
|
6 2013-12-24 2013-12-24 65 2 NaN
|
||||||
|
7 2013-12-24 2013-12-24 66 3 3
|
||||||
|
8 2013-12-24 2013-12-24 67 NaN 4
|
||||||
|
|
||||||
|
output (expected):
|
||||||
|
other value
|
||||||
|
2014-01-01 Equity(65 [A]) 2 1
|
||||||
|
Equity(66 [B]) 3 3
|
||||||
|
Equity(67 [C]) 3 4
|
||||||
|
2014-01-02 Equity(65 [A]) 2 1
|
||||||
|
Equity(66 [B]) 3 3
|
||||||
|
Equity(67 [C]) 3 4
|
||||||
|
2014-01-03 Equity(65 [A]) 2 1
|
||||||
|
Equity(66 [B]) 3 3
|
||||||
|
Equity(67 [C]) 3 4
|
||||||
|
"""
|
||||||
|
dates = self.dates.repeat(3) - timedelta(days=10)
|
||||||
|
df = pd.DataFrame({
|
||||||
|
'sid': self.sids * 3,
|
||||||
|
'value': (0, 1, np.nan, 1, np.nan, 3, np.nan, 3, 4),
|
||||||
|
'other': (0, np.nan, 2, np.nan, 2, 3, 2, 3, np.nan),
|
||||||
|
'asof_date': dates,
|
||||||
|
'timestamp': dates,
|
||||||
|
})
|
||||||
|
fields = OrderedDict(self.dshape.measure.fields)
|
||||||
|
fields['other'] = fields['value']
|
||||||
|
|
||||||
|
with tmp_asset_finder() as finder:
|
||||||
|
expected = pd.DataFrame(
|
||||||
|
np.array([[2, 1],
|
||||||
|
[3, 3],
|
||||||
|
[3, 4],
|
||||||
|
[2, 1],
|
||||||
|
[3, 3],
|
||||||
|
[3, 4],
|
||||||
|
[2, 1],
|
||||||
|
[3, 3],
|
||||||
|
[3, 4]]),
|
||||||
|
columns=['other', 'value'],
|
||||||
|
index=pd.MultiIndex.from_product(
|
||||||
|
(self.dates, finder.retrieve_all(self.sids)),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
self._test_id(
|
||||||
|
df,
|
||||||
|
var * Record(fields),
|
||||||
|
expected,
|
||||||
|
finder,
|
||||||
|
('value', 'other'),
|
||||||
|
)
|
||||||
|
|
||||||
def test_id_multiple_columns(self):
|
def test_id_multiple_columns(self):
|
||||||
|
"""
|
||||||
|
input (df):
|
||||||
|
asof_date sid timestamp value other
|
||||||
|
0 2014-01-01 65 2014-01-01 0 1
|
||||||
|
1 2014-01-01 66 2014-01-01 1 2
|
||||||
|
2 2014-01-01 67 2014-01-01 2 3
|
||||||
|
3 2014-01-02 65 2014-01-02 1 2
|
||||||
|
4 2014-01-02 66 2014-01-02 2 3
|
||||||
|
5 2014-01-02 67 2014-01-02 3 4
|
||||||
|
6 2014-01-03 65 2014-01-03 2 3
|
||||||
|
7 2014-01-03 66 2014-01-03 3 4
|
||||||
|
8 2014-01-03 67 2014-01-03 4 5
|
||||||
|
|
||||||
|
output (expected):
|
||||||
|
value other
|
||||||
|
2014-01-01 Equity(65 [A]) 0 1
|
||||||
|
Equity(66 [B]) 1 2
|
||||||
|
Equity(67 [C]) 2 3
|
||||||
|
2014-01-02 Equity(65 [A]) 1 2
|
||||||
|
Equity(66 [B]) 2 3
|
||||||
|
Equity(67 [C]) 3 4
|
||||||
|
2014-01-03 Equity(65 [A]) 2 3
|
||||||
|
Equity(66 [B]) 3 4
|
||||||
|
Equity(67 [C]) 4 5
|
||||||
|
"""
|
||||||
df = self.df.copy()
|
df = self.df.copy()
|
||||||
df['other'] = df.value + 1
|
df['other'] = df.value + 1
|
||||||
fields = OrderedDict(self.dshape.measure.fields)
|
fields = OrderedDict(self.dshape.measure.fields)
|
||||||
fields['other'] = fields['value']
|
fields['other'] = fields['value']
|
||||||
expr = bz.Data(df, name='expr', dshape=var * Record(fields))
|
|
||||||
loader = BlazeLoader()
|
|
||||||
ds = from_blaze(
|
|
||||||
expr,
|
|
||||||
loader=loader,
|
|
||||||
no_deltas_rule='ignore',
|
|
||||||
)
|
|
||||||
p = Pipeline()
|
|
||||||
p.add(ds.value.latest, 'value')
|
|
||||||
p.add(ds.other.latest, 'other')
|
|
||||||
dates = self.dates
|
|
||||||
|
|
||||||
with tmp_asset_finder() as finder:
|
with tmp_asset_finder() as finder:
|
||||||
result = SimplePipelineEngine(
|
expected = df.drop('asof_date', axis=1).set_index(
|
||||||
loader,
|
['timestamp', 'sid'],
|
||||||
dates,
|
).sort_index(axis=1)
|
||||||
|
expected.index = pd.MultiIndex.from_product((
|
||||||
|
expected.index.levels[0],
|
||||||
|
finder.retrieve_all(expected.index.levels[1]),
|
||||||
|
))
|
||||||
|
self._test_id(
|
||||||
|
df,
|
||||||
|
var * Record(fields),
|
||||||
|
expected,
|
||||||
finder,
|
finder,
|
||||||
).run_pipeline(p, dates[0], dates[-1])
|
('value', 'other'),
|
||||||
|
)
|
||||||
expected = df.drop('asof_date', axis=1).set_index(
|
|
||||||
['timestamp', 'sid'],
|
|
||||||
).sort_index(axis=1)
|
|
||||||
expected.index = pd.MultiIndex.from_product((
|
|
||||||
expected.index.levels[0],
|
|
||||||
finder.retrieve_all(expected.index.levels[1]),
|
|
||||||
))
|
|
||||||
assert_frame_equal(
|
|
||||||
result,
|
|
||||||
expected.sort_index(axis=1),
|
|
||||||
check_dtype=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_id_macro_dataset(self):
|
def test_id_macro_dataset(self):
|
||||||
expr = bz.Data(self.macro_df, name='expr', dshape=self.macro_dshape)
|
"""
|
||||||
loader = BlazeLoader()
|
input (self.macro_df)
|
||||||
ds = from_blaze(
|
asof_date timestamp value
|
||||||
expr,
|
0 2014-01-01 2014-01-01 0
|
||||||
loader=loader,
|
3 2014-01-02 2014-01-02 1
|
||||||
no_deltas_rule=no_deltas_rules.ignore,
|
6 2014-01-03 2014-01-03 2
|
||||||
)
|
|
||||||
p = Pipeline()
|
|
||||||
p.add(ds.value.latest, 'value')
|
|
||||||
dates = self.dates
|
|
||||||
|
|
||||||
|
output (expected):
|
||||||
|
value
|
||||||
|
2014-01-01 Equity(65 [A]) 0
|
||||||
|
Equity(66 [B]) 0
|
||||||
|
Equity(67 [C]) 0
|
||||||
|
2014-01-02 Equity(65 [A]) 1
|
||||||
|
Equity(66 [B]) 1
|
||||||
|
Equity(67 [C]) 1
|
||||||
|
2014-01-03 Equity(65 [A]) 2
|
||||||
|
Equity(66 [B]) 2
|
||||||
|
Equity(67 [C]) 2
|
||||||
|
"""
|
||||||
asset_info = asset_infos[0][0]
|
asset_info = asset_infos[0][0]
|
||||||
with tmp_asset_finder(equities=asset_info) as finder:
|
|
||||||
result = SimplePipelineEngine(
|
|
||||||
loader,
|
|
||||||
dates,
|
|
||||||
finder,
|
|
||||||
).run_pipeline(p, dates[0], dates[-1])
|
|
||||||
|
|
||||||
nassets = len(asset_info)
|
nassets = len(asset_info)
|
||||||
expected = pd.DataFrame(
|
with tmp_asset_finder() as finder:
|
||||||
list(concatv([0] * nassets, [1] * nassets, [2] * nassets)),
|
expected = pd.DataFrame(
|
||||||
index=pd.MultiIndex.from_product((
|
list(concatv([0] * nassets, [1] * nassets, [2] * nassets)),
|
||||||
self.macro_df.timestamp,
|
index=pd.MultiIndex.from_product((
|
||||||
finder.retrieve_all(asset_info.index),
|
self.macro_df.timestamp,
|
||||||
)),
|
finder.retrieve_all(asset_info.index),
|
||||||
columns=('value',),
|
)),
|
||||||
)
|
columns=('value',),
|
||||||
assert_frame_equal(result, expected, check_dtype=False)
|
)
|
||||||
|
self._test_id(
|
||||||
|
self.macro_df,
|
||||||
|
self.macro_dshape,
|
||||||
|
expected,
|
||||||
|
finder,
|
||||||
|
('value',),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_id_ffill_out_of_window_macro_dataset(self):
|
||||||
|
"""
|
||||||
|
input (df):
|
||||||
|
asof_date timestamp other value
|
||||||
|
0 2013-12-22 2013-12-22 NaN 0
|
||||||
|
1 2013-12-23 2013-12-23 1 NaN
|
||||||
|
2 2013-12-24 2013-12-24 NaN NaN
|
||||||
|
|
||||||
|
output (expected):
|
||||||
|
other value
|
||||||
|
2014-01-01 Equity(65 [A]) 1 0
|
||||||
|
Equity(66 [B]) 1 0
|
||||||
|
Equity(67 [C]) 1 0
|
||||||
|
2014-01-02 Equity(65 [A]) 1 0
|
||||||
|
Equity(66 [B]) 1 0
|
||||||
|
Equity(67 [C]) 1 0
|
||||||
|
2014-01-03 Equity(65 [A]) 1 0
|
||||||
|
Equity(66 [B]) 1 0
|
||||||
|
Equity(67 [C]) 1 0
|
||||||
|
"""
|
||||||
|
dates = self.dates - timedelta(days=10)
|
||||||
|
df = pd.DataFrame({
|
||||||
|
'value': (0, np.nan, np.nan),
|
||||||
|
'other': (np.nan, 1, np.nan),
|
||||||
|
'asof_date': dates,
|
||||||
|
'timestamp': dates,
|
||||||
|
})
|
||||||
|
fields = OrderedDict(self.macro_dshape.measure.fields)
|
||||||
|
fields['other'] = fields['value']
|
||||||
|
|
||||||
|
with tmp_asset_finder() as finder:
|
||||||
|
expected = pd.DataFrame(
|
||||||
|
np.array([[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1],
|
||||||
|
[0, 1]]),
|
||||||
|
columns=['value', 'other'],
|
||||||
|
index=pd.MultiIndex.from_product(
|
||||||
|
(self.dates, finder.retrieve_all(self.sids)),
|
||||||
|
),
|
||||||
|
).sort_index(axis=1)
|
||||||
|
self._test_id(
|
||||||
|
df,
|
||||||
|
var * Record(fields),
|
||||||
|
expected,
|
||||||
|
finder,
|
||||||
|
('value', 'other'),
|
||||||
|
)
|
||||||
|
|
||||||
def test_id_macro_dataset_multiple_columns(self):
|
def test_id_macro_dataset_multiple_columns(self):
|
||||||
|
"""
|
||||||
|
input (df):
|
||||||
|
asof_date timestamp other value
|
||||||
|
0 2014-01-01 2014-01-01 1 0
|
||||||
|
3 2014-01-02 2014-01-02 2 1
|
||||||
|
6 2014-01-03 2014-01-03 3 2
|
||||||
|
|
||||||
|
output (expected):
|
||||||
|
other value
|
||||||
|
2014-01-01 Equity(65 [A]) 1 0
|
||||||
|
Equity(66 [B]) 1 0
|
||||||
|
Equity(67 [C]) 1 0
|
||||||
|
2014-01-02 Equity(65 [A]) 2 1
|
||||||
|
Equity(66 [B]) 2 1
|
||||||
|
Equity(67 [C]) 2 1
|
||||||
|
2014-01-03 Equity(65 [A]) 3 2
|
||||||
|
Equity(66 [B]) 3 2
|
||||||
|
Equity(67 [C]) 3 2
|
||||||
|
"""
|
||||||
df = self.macro_df.copy()
|
df = self.macro_df.copy()
|
||||||
df['other'] = df.value + 1
|
df['other'] = df.value + 1
|
||||||
fields = OrderedDict(self.macro_dshape.measure.fields)
|
fields = OrderedDict(self.macro_dshape.measure.fields)
|
||||||
@@ -438,15 +620,15 @@ class BlazeToPipelineTestCase(TestCase):
|
|||||||
ds = from_blaze(
|
ds = from_blaze(
|
||||||
expr,
|
expr,
|
||||||
loader=loader,
|
loader=loader,
|
||||||
no_deltas_rule='ignore',
|
no_deltas_rule=no_deltas_rules.ignore,
|
||||||
)
|
)
|
||||||
p = Pipeline()
|
p = Pipeline()
|
||||||
p.add(ds.value.latest, 'value')
|
p.add(ds.value.latest, 'value')
|
||||||
p.add(ds.other.latest, 'latest')
|
p.add(ds.other.latest, 'other')
|
||||||
dates = self.dates
|
dates = self.dates
|
||||||
|
|
||||||
asset_info = asset_infos[0][0]
|
asset_info = asset_infos[0][0]
|
||||||
with tmp_asset_finder(asset_info) as finder:
|
with tmp_asset_finder(equities=asset_info) as finder:
|
||||||
result = SimplePipelineEngine(
|
result = SimplePipelineEngine(
|
||||||
loader,
|
loader,
|
||||||
dates,
|
dates,
|
||||||
@@ -461,7 +643,7 @@ class BlazeToPipelineTestCase(TestCase):
|
|||||||
df.timestamp,
|
df.timestamp,
|
||||||
finder.retrieve_all(asset_info.index),
|
finder.retrieve_all(asset_info.index),
|
||||||
)),
|
)),
|
||||||
columns=('value', 'latest'),
|
columns=('value', 'other'),
|
||||||
).sort_index(axis=1)
|
).sort_index(axis=1)
|
||||||
assert_frame_equal(
|
assert_frame_equal(
|
||||||
result,
|
result,
|
||||||
|
|||||||
Reference in New Issue
Block a user