mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-08 09:33:24 +08:00
llllllllll
918 lines
30 KiB
Python
918 lines
30 KiB
Python
"""
|
|
Tests for the blaze interface to the pipeline api.
|
|
"""
|
|
from __future__ import division
|
|
|
|
from collections import OrderedDict
|
|
from datetime import timedelta, time
|
|
from unittest import TestCase
|
|
import warnings
|
|
|
|
import blaze as bz
|
|
from datashape import dshape, var, Record
|
|
from nose_parameterized import parameterized
|
|
import numpy as np
|
|
from numpy.testing.utils import assert_array_almost_equal
|
|
import pandas as pd
|
|
from pandas.util.testing import assert_frame_equal
|
|
from toolz import keymap, valmap, concatv
|
|
from toolz.curried import operator as op
|
|
|
|
from zipline.pipeline import Pipeline, CustomFactor
|
|
from zipline.pipeline.data import DataSet, BoundColumn
|
|
from zipline.pipeline.engine import SimplePipelineEngine
|
|
from zipline.pipeline.loaders.blaze import (
|
|
from_blaze,
|
|
BlazeLoader,
|
|
NoDeltasWarning,
|
|
)
|
|
from zipline.pipeline.loaders.blaze.core import (
|
|
NonNumpyField,
|
|
NonPipelineField,
|
|
no_deltas_rules,
|
|
)
|
|
from zipline.utils.numpy_utils import repeat_last_axis
|
|
from zipline.utils.test_utils import tmp_asset_finder, make_simple_equity_info
|
|
|
|
|
|
nameof = op.attrgetter('name')
|
|
dtypeof = op.attrgetter('dtype')
|
|
asset_infos = (
|
|
(make_simple_equity_info(
|
|
tuple(map(ord, 'ABC')),
|
|
pd.Timestamp(0),
|
|
pd.Timestamp('2015'),
|
|
),),
|
|
(make_simple_equity_info(
|
|
tuple(map(ord, 'ABCD')),
|
|
pd.Timestamp(0),
|
|
pd.Timestamp('2015'),
|
|
),),
|
|
)
|
|
with_extra_sid = parameterized.expand(asset_infos)
|
|
|
|
|
|
class BlazeToPipelineTestCase(TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.dates = dates = pd.date_range('2014-01-01', '2014-01-03')
|
|
dates = cls.dates.repeat(3)
|
|
cls.sids = sids = ord('A'), ord('B'), ord('C')
|
|
cls.df = df = pd.DataFrame({
|
|
'sid': sids * 3,
|
|
'value': (0, 1, 2, 1, 2, 3, 2, 3, 4),
|
|
'asof_date': dates,
|
|
'timestamp': dates,
|
|
})
|
|
cls.dshape = dshape("""
|
|
var * {
|
|
sid: ?int64,
|
|
value: ?float64,
|
|
asof_date: datetime,
|
|
timestamp: datetime
|
|
}
|
|
""")
|
|
cls.macro_df = df[df.sid == 65].drop('sid', axis=1)
|
|
dshape_ = OrderedDict(cls.dshape.measure.fields)
|
|
del dshape_['sid']
|
|
cls.macro_dshape = var * Record(dshape_)
|
|
|
|
cls.garbage_loader = BlazeLoader()
|
|
|
|
def test_tabular(self):
|
|
name = 'expr'
|
|
expr = bz.Data(self.df, name=name, dshape=self.dshape)
|
|
ds = from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
self.assertEqual(ds.__name__, name)
|
|
self.assertTrue(issubclass(ds, DataSet))
|
|
self.assertEqual(
|
|
{c.name: c.dtype for c in ds.columns},
|
|
{'sid': np.int64, 'value': np.float64},
|
|
)
|
|
|
|
for field in ('timestamp', 'asof_date'):
|
|
with self.assertRaises(AttributeError) as e:
|
|
getattr(ds, field)
|
|
self.assertIn("'%s'" % field, str(e.exception))
|
|
self.assertIn("'datetime'", str(e.exception))
|
|
|
|
# test memoization
|
|
self.assertIs(
|
|
from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
),
|
|
ds,
|
|
)
|
|
|
|
def test_column(self):
|
|
exprname = 'expr'
|
|
expr = bz.Data(self.df, name=exprname, dshape=self.dshape)
|
|
value = from_blaze(
|
|
expr.value,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
self.assertEqual(value.name, 'value')
|
|
self.assertIsInstance(value, BoundColumn)
|
|
self.assertEqual(value.dtype, np.float64)
|
|
|
|
# test memoization
|
|
self.assertIs(
|
|
from_blaze(
|
|
expr.value,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
),
|
|
value,
|
|
)
|
|
self.assertIs(
|
|
from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
).value,
|
|
value,
|
|
)
|
|
|
|
# test the walk back up the tree
|
|
self.assertIs(
|
|
from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
),
|
|
value.dataset,
|
|
)
|
|
self.assertEqual(value.dataset.__name__, exprname)
|
|
|
|
def test_missing_asof(self):
|
|
expr = bz.Data(
|
|
self.df.loc[:, ['sid', 'value', 'timestamp']],
|
|
name='expr',
|
|
dshape="""
|
|
var * {
|
|
sid: ?int64,
|
|
value: float64,
|
|
timestamp: datetime,
|
|
}""",
|
|
)
|
|
|
|
with self.assertRaises(TypeError) as e:
|
|
from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
self.assertIn("'asof_date'", str(e.exception))
|
|
self.assertIn(repr(str(expr.dshape.measure)), str(e.exception))
|
|
|
|
def test_auto_deltas(self):
|
|
expr = bz.Data(
|
|
{'ds': self.df,
|
|
'ds_deltas': pd.DataFrame(columns=self.df.columns)},
|
|
dshape=var * Record((
|
|
('ds', self.dshape.measure),
|
|
('ds_deltas', self.dshape.measure),
|
|
)),
|
|
)
|
|
loader = BlazeLoader()
|
|
ds = from_blaze(expr.ds, loader=loader)
|
|
self.assertEqual(len(loader), 1)
|
|
exprdata = loader[ds]
|
|
self.assertTrue(exprdata.expr.isidentical(expr.ds))
|
|
self.assertTrue(exprdata.deltas.isidentical(expr.ds_deltas))
|
|
|
|
def test_auto_deltas_fail_warn(self):
|
|
with warnings.catch_warnings(record=True) as ws:
|
|
warnings.simplefilter('always')
|
|
loader = BlazeLoader()
|
|
expr = bz.Data(self.df, dshape=self.dshape)
|
|
from_blaze(
|
|
expr,
|
|
loader=loader,
|
|
no_deltas_rule=no_deltas_rules.warn,
|
|
)
|
|
self.assertEqual(len(ws), 1)
|
|
w = ws[0].message
|
|
self.assertIsInstance(w, NoDeltasWarning)
|
|
self.assertIn(str(expr), str(w))
|
|
|
|
def test_auto_deltas_fail_raise(self):
|
|
loader = BlazeLoader()
|
|
expr = bz.Data(self.df, dshape=self.dshape)
|
|
with self.assertRaises(ValueError) as e:
|
|
from_blaze(
|
|
expr,
|
|
loader=loader,
|
|
no_deltas_rule=no_deltas_rules.raise_,
|
|
)
|
|
self.assertIn(str(expr), str(e.exception))
|
|
|
|
def test_non_numpy_field(self):
|
|
expr = bz.Data(
|
|
[],
|
|
dshape="""
|
|
var * {
|
|
a: datetime,
|
|
asof_date: datetime,
|
|
timestamp: datetime,
|
|
}""",
|
|
)
|
|
ds = from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
with self.assertRaises(AttributeError):
|
|
ds.a
|
|
self.assertIsInstance(object.__getattribute__(ds, 'a'), NonNumpyField)
|
|
|
|
def test_non_pipeline_field(self):
|
|
# NOTE: This test will fail if we ever allow string types in
|
|
# the Pipeline API. If this happens, change the dtype of the `a` field
|
|
# of expr to another type we don't allow.
|
|
expr = bz.Data(
|
|
[],
|
|
dshape="""
|
|
var * {
|
|
a: string,
|
|
asof_date: datetime,
|
|
timestamp: datetime,
|
|
}""",
|
|
)
|
|
ds = from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
with self.assertRaises(AttributeError):
|
|
ds.a
|
|
self.assertIsInstance(
|
|
object.__getattribute__(ds, 'a'),
|
|
NonPipelineField,
|
|
)
|
|
|
|
def test_complex_expr(self):
|
|
expr = bz.Data(self.df, dshape=self.dshape)
|
|
# put an Add in the table
|
|
expr_with_add = bz.transform(expr, value=expr.value + 1)
|
|
|
|
# Test that we can have complex expressions with no deltas
|
|
from_blaze(
|
|
expr_with_add,
|
|
deltas=None,
|
|
loader=self.garbage_loader,
|
|
)
|
|
|
|
with self.assertRaises(TypeError):
|
|
from_blaze(
|
|
expr.value + 1, # put an Add in the column
|
|
deltas=None,
|
|
loader=self.garbage_loader,
|
|
)
|
|
|
|
deltas = bz.Data(
|
|
pd.DataFrame(columns=self.df.columns),
|
|
dshape=self.dshape,
|
|
)
|
|
with self.assertRaises(TypeError):
|
|
from_blaze(
|
|
expr_with_add,
|
|
deltas=deltas,
|
|
loader=self.garbage_loader,
|
|
)
|
|
|
|
with self.assertRaises(TypeError):
|
|
from_blaze(
|
|
expr.value + 1,
|
|
deltas=deltas,
|
|
loader=self.garbage_loader,
|
|
)
|
|
|
|
def _test_id(self, df, dshape, expected, finder, add):
|
|
expr = bz.Data(df, name='expr', dshape=dshape)
|
|
loader = BlazeLoader()
|
|
ds = from_blaze(
|
|
expr,
|
|
loader=loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
p = Pipeline()
|
|
for a in add:
|
|
p.add(getattr(ds, a).latest, a)
|
|
dates = self.dates
|
|
|
|
with tmp_asset_finder() as finder:
|
|
result = SimplePipelineEngine(
|
|
loader,
|
|
dates,
|
|
finder,
|
|
).run_pipeline(p, dates[0], dates[-1])
|
|
|
|
assert_frame_equal(result, expected, check_dtype=False)
|
|
|
|
def test_custom_query_time_tz(self):
|
|
df = self.df.copy()
|
|
df['timestamp'] = (
|
|
pd.DatetimeIndex(df['timestamp'], tz='EST') +
|
|
timedelta(hours=8, minutes=44)
|
|
).tz_convert('utc').tz_localize(None)
|
|
df.ix[3:5, 'timestamp'] = pd.Timestamp('2014-01-01 13:45')
|
|
expr = bz.Data(df, name='expr', dshape=self.dshape)
|
|
loader = BlazeLoader(data_query_time=time(8, 45), data_query_tz='EST')
|
|
ds = from_blaze(
|
|
expr,
|
|
loader=loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
)
|
|
p = Pipeline()
|
|
p.add(ds.value.latest, 'value')
|
|
dates = self.dates
|
|
|
|
with tmp_asset_finder() as finder:
|
|
result = SimplePipelineEngine(
|
|
loader,
|
|
dates,
|
|
finder,
|
|
).run_pipeline(p, dates[0], dates[-1])
|
|
|
|
expected = df.drop('asof_date', axis=1)
|
|
expected['timestamp'] = expected['timestamp'].dt.normalize().astype(
|
|
'datetime64[ns]',
|
|
)
|
|
expected.ix[3:5, 'timestamp'] += timedelta(days=1)
|
|
expected.set_index(['timestamp', 'sid'], inplace=True)
|
|
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)
|
|
|
|
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):
|
|
"""
|
|
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['other'] = df.value + 1
|
|
fields = OrderedDict(self.dshape.measure.fields)
|
|
fields['other'] = fields['value']
|
|
with tmp_asset_finder() as finder:
|
|
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]),
|
|
))
|
|
self._test_id(
|
|
df,
|
|
var * Record(fields),
|
|
expected,
|
|
finder,
|
|
('value', 'other'),
|
|
)
|
|
|
|
def test_id_macro_dataset(self):
|
|
"""
|
|
input (self.macro_df)
|
|
asof_date timestamp value
|
|
0 2014-01-01 2014-01-01 0
|
|
3 2014-01-02 2014-01-02 1
|
|
6 2014-01-03 2014-01-03 2
|
|
|
|
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]
|
|
nassets = len(asset_info)
|
|
with tmp_asset_finder() as finder:
|
|
expected = pd.DataFrame(
|
|
list(concatv([0] * nassets, [1] * nassets, [2] * nassets)),
|
|
index=pd.MultiIndex.from_product((
|
|
self.macro_df.timestamp,
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
columns=('value',),
|
|
)
|
|
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):
|
|
"""
|
|
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['other'] = df.value + 1
|
|
fields = OrderedDict(self.macro_dshape.measure.fields)
|
|
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=no_deltas_rules.ignore,
|
|
)
|
|
p = Pipeline()
|
|
p.add(ds.value.latest, 'value')
|
|
p.add(ds.other.latest, 'other')
|
|
dates = self.dates
|
|
|
|
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])
|
|
|
|
expected = pd.DataFrame(
|
|
np.array([[0, 1],
|
|
[1, 2],
|
|
[2, 3]]).repeat(3, axis=0),
|
|
index=pd.MultiIndex.from_product((
|
|
df.timestamp,
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
columns=('value', 'other'),
|
|
).sort_index(axis=1)
|
|
assert_frame_equal(
|
|
result,
|
|
expected.sort_index(axis=1),
|
|
check_dtype=False,
|
|
)
|
|
|
|
def _run_pipeline(self,
|
|
expr,
|
|
deltas,
|
|
expected_views,
|
|
expected_output,
|
|
finder,
|
|
calendar,
|
|
start,
|
|
end,
|
|
window_length,
|
|
compute_fn):
|
|
loader = BlazeLoader()
|
|
ds = from_blaze(
|
|
expr,
|
|
deltas,
|
|
loader=loader,
|
|
no_deltas_rule=no_deltas_rules.raise_,
|
|
)
|
|
p = Pipeline()
|
|
|
|
# prevent unbound locals issue in the inner class
|
|
window_length_ = window_length
|
|
|
|
class TestFactor(CustomFactor):
|
|
inputs = ds.value,
|
|
window_length = window_length_
|
|
|
|
def compute(self, today, assets, out, data):
|
|
assert_array_almost_equal(data, expected_views[today])
|
|
out[:] = compute_fn(data)
|
|
|
|
p.add(TestFactor(), 'value')
|
|
|
|
result = SimplePipelineEngine(
|
|
loader,
|
|
calendar,
|
|
finder,
|
|
).run_pipeline(p, start, end)
|
|
|
|
assert_frame_equal(
|
|
result,
|
|
expected_output,
|
|
check_dtype=False,
|
|
)
|
|
|
|
@with_extra_sid
|
|
def test_deltas(self, asset_info):
|
|
expr = bz.Data(self.df, name='expr', dshape=self.dshape)
|
|
deltas = bz.Data(self.df, name='deltas', dshape=self.dshape)
|
|
deltas = bz.transform(
|
|
deltas,
|
|
value=deltas.value + 10,
|
|
timestamp=deltas.timestamp + timedelta(days=1),
|
|
)
|
|
|
|
expected_views = keymap(pd.Timestamp, {
|
|
'2014-01-02': np.array([[10.0, 11.0, 12.0],
|
|
[1.0, 2.0, 3.0]]),
|
|
'2014-01-03': np.array([[11.0, 12.0, 13.0],
|
|
[2.0, 3.0, 4.0]]),
|
|
'2014-01-04': np.array([[12.0, 13.0, 14.0],
|
|
[12.0, 13.0, 14.0]]),
|
|
})
|
|
|
|
nassets = len(asset_info)
|
|
if nassets == 4:
|
|
expected_views = valmap(
|
|
lambda view: np.c_[view, [np.nan, np.nan]],
|
|
expected_views,
|
|
)
|
|
|
|
with tmp_asset_finder(equities=asset_info) as finder:
|
|
expected_output = pd.DataFrame(
|
|
list(concatv([12] * nassets, [13] * nassets, [14] * nassets)),
|
|
index=pd.MultiIndex.from_product((
|
|
sorted(expected_views.keys()),
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
columns=('value',),
|
|
)
|
|
dates = self.dates
|
|
dates = dates.insert(len(dates), dates[-1] + timedelta(days=1))
|
|
self._run_pipeline(
|
|
expr,
|
|
deltas,
|
|
expected_views,
|
|
expected_output,
|
|
finder,
|
|
calendar=dates,
|
|
start=dates[1],
|
|
end=dates[-1],
|
|
window_length=2,
|
|
compute_fn=np.nanmax,
|
|
)
|
|
|
|
def test_deltas_macro(self):
|
|
asset_info = asset_infos[0][0]
|
|
expr = bz.Data(self.macro_df, name='expr', dshape=self.macro_dshape)
|
|
deltas = bz.Data(
|
|
self.macro_df.iloc[:-1],
|
|
name='deltas',
|
|
dshape=self.macro_dshape,
|
|
)
|
|
deltas = bz.transform(
|
|
deltas,
|
|
value=deltas.value + 10,
|
|
timestamp=deltas.timestamp + timedelta(days=1),
|
|
)
|
|
|
|
nassets = len(asset_info)
|
|
expected_views = keymap(pd.Timestamp, {
|
|
'2014-01-02': repeat_last_axis(np.array([10.0, 1.0]), nassets),
|
|
'2014-01-03': repeat_last_axis(np.array([11.0, 2.0]), nassets),
|
|
})
|
|
|
|
with tmp_asset_finder(equities=asset_info) as finder:
|
|
expected_output = pd.DataFrame(
|
|
list(concatv([10] * nassets, [11] * nassets)),
|
|
index=pd.MultiIndex.from_product((
|
|
sorted(expected_views.keys()),
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
columns=('value',),
|
|
)
|
|
dates = self.dates
|
|
self._run_pipeline(
|
|
expr,
|
|
deltas,
|
|
expected_views,
|
|
expected_output,
|
|
finder,
|
|
calendar=dates,
|
|
start=dates[1],
|
|
end=dates[-1],
|
|
window_length=2,
|
|
compute_fn=np.nanmax,
|
|
)
|
|
|
|
@with_extra_sid
|
|
def test_novel_deltas(self, asset_info):
|
|
base_dates = pd.DatetimeIndex([
|
|
pd.Timestamp('2014-01-01'),
|
|
pd.Timestamp('2014-01-04')
|
|
])
|
|
repeated_dates = base_dates.repeat(3)
|
|
baseline = pd.DataFrame({
|
|
'sid': self.sids * 2,
|
|
'value': (0, 1, 2, 1, 2, 3),
|
|
'asof_date': repeated_dates,
|
|
'timestamp': repeated_dates,
|
|
})
|
|
expr = bz.Data(baseline, name='expr', dshape=self.dshape)
|
|
deltas = bz.Data(baseline, name='deltas', dshape=self.dshape)
|
|
deltas = bz.transform(
|
|
deltas,
|
|
value=deltas.value + 10,
|
|
timestamp=deltas.timestamp + timedelta(days=1),
|
|
)
|
|
expected_views = keymap(pd.Timestamp, {
|
|
'2014-01-03': np.array([[10.0, 11.0, 12.0],
|
|
[10.0, 11.0, 12.0],
|
|
[10.0, 11.0, 12.0]]),
|
|
'2014-01-06': np.array([[10.0, 11.0, 12.0],
|
|
[10.0, 11.0, 12.0],
|
|
[11.0, 12.0, 13.0]]),
|
|
})
|
|
if len(asset_info) == 4:
|
|
expected_views = valmap(
|
|
lambda view: np.c_[view, [np.nan, np.nan, np.nan]],
|
|
expected_views,
|
|
)
|
|
expected_output_buffer = [10, 11, 12, np.nan, 11, 12, 13, np.nan]
|
|
else:
|
|
expected_output_buffer = [10, 11, 12, 11, 12, 13]
|
|
|
|
cal = pd.DatetimeIndex([
|
|
pd.Timestamp('2014-01-01'),
|
|
pd.Timestamp('2014-01-02'),
|
|
pd.Timestamp('2014-01-03'),
|
|
# omitting the 4th and 5th to simulate a weekend
|
|
pd.Timestamp('2014-01-06'),
|
|
])
|
|
|
|
with tmp_asset_finder(equities=asset_info) as finder:
|
|
expected_output = pd.DataFrame(
|
|
expected_output_buffer,
|
|
index=pd.MultiIndex.from_product((
|
|
sorted(expected_views.keys()),
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
columns=('value',),
|
|
)
|
|
self._run_pipeline(
|
|
expr,
|
|
deltas,
|
|
expected_views,
|
|
expected_output,
|
|
finder,
|
|
calendar=cal,
|
|
start=cal[2],
|
|
end=cal[-1],
|
|
window_length=3,
|
|
compute_fn=op.itemgetter(-1),
|
|
)
|
|
|
|
def test_novel_deltas_macro(self):
|
|
asset_info = asset_infos[0][0]
|
|
base_dates = pd.DatetimeIndex([
|
|
pd.Timestamp('2014-01-01'),
|
|
pd.Timestamp('2014-01-04')
|
|
])
|
|
baseline = pd.DataFrame({
|
|
'value': (0, 1),
|
|
'asof_date': base_dates,
|
|
'timestamp': base_dates,
|
|
})
|
|
expr = bz.Data(baseline, name='expr', dshape=self.macro_dshape)
|
|
deltas = bz.Data(baseline, name='deltas', dshape=self.macro_dshape)
|
|
deltas = bz.transform(
|
|
deltas,
|
|
value=deltas.value + 10,
|
|
timestamp=deltas.timestamp + timedelta(days=1),
|
|
)
|
|
|
|
nassets = len(asset_info)
|
|
expected_views = keymap(pd.Timestamp, {
|
|
'2014-01-03': repeat_last_axis(
|
|
np.array([10.0, 10.0, 10.0]),
|
|
nassets,
|
|
),
|
|
'2014-01-06': repeat_last_axis(
|
|
np.array([10.0, 10.0, 11.0]),
|
|
nassets,
|
|
),
|
|
})
|
|
|
|
cal = pd.DatetimeIndex([
|
|
pd.Timestamp('2014-01-01'),
|
|
pd.Timestamp('2014-01-02'),
|
|
pd.Timestamp('2014-01-03'),
|
|
# omitting the 4th and 5th to simulate a weekend
|
|
pd.Timestamp('2014-01-06'),
|
|
])
|
|
with tmp_asset_finder(equities=asset_info) as finder:
|
|
expected_output = pd.DataFrame(
|
|
list(concatv([10] * nassets, [11] * nassets)),
|
|
index=pd.MultiIndex.from_product((
|
|
sorted(expected_views.keys()),
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
columns=('value',),
|
|
)
|
|
self._run_pipeline(
|
|
expr,
|
|
deltas,
|
|
expected_views,
|
|
expected_output,
|
|
finder,
|
|
calendar=cal,
|
|
start=cal[2],
|
|
end=cal[-1],
|
|
window_length=3,
|
|
compute_fn=op.itemgetter(-1),
|
|
)
|