mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-08 05:36:44 +08:00
Joe Jevnik
bc0b117dc9
Changes BcolzDailyBarWriter to not be an abc, data is passed as an iterator of (sid, dataframe) pairs to the write method. Changes the AssetsDBWriter to be a single class which accepts an engine at construction time and has a `write` method for writing dataframes for the various tables. We no longer support writing the various other data types, callers should coerce their data into a dataframe themselves. See zipline.assets.synthetic for some helpers to do this. Adds many new fixtures and updates some existing fixtures to use the new ones: WithDefaultDateBounds A fixture that provides the suite a START_DATE and END_DATE. This is meant to make it easy for other fixtures to synchronize their date ranges without depending on eachother in strange ways. For example, WithBcolzMinuteBarReader and WithBcolzDailyBarReader by default should both have data for the same dates, so they may use depend on WithDefaultDates without forcing a dependency between them. WithTmpDir, WithInstanceTmpDir Provides the suite or individual test case a temporary directory. WithBcolzDailyBarReader Provides the suite a BcolzDailyBarReader which reads from bcolz data written to a temporary directory. The data will be read from dataframes and then converted to bcolz files with BcolzDailyBarWriter.write WithBcolzDailyBarReaderFromCSVs Provides the suite a BcolzDailyBarReader which reads from bcolz data written to a temporary directory. The data will be read from a collection of CSV files and then converted into the bcolz data through BcolzDailyBarWriter.write_csvs WithBcolzMinuteBarReader Provides the suite a BcolzMinuteBarReader which reads from bcolz data written to a temporary directory. The data will be read from dataframes and then converted to bcolz files with BcolzMinuteBarWriter.write WithAdjustmentReader Provides the suite a SQLiteAdjustmentReader which reads from an in memory sqlite database. The data will be read from dataframes and then converted into sqlite with SQLiteAdjustmentWriter.write WithDataPortal Provides each test case a DataPortal object with data from temporary resources.
1143 lines
39 KiB
Python
1143 lines
39 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 itertools import product, chain
|
|
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
|
|
from odo import odo
|
|
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.assets.synthetic import make_simple_equity_info
|
|
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 (
|
|
float64_dtype,
|
|
int64_dtype,
|
|
repeat_last_axis,
|
|
)
|
|
from zipline.testing import tmp_asset_finder
|
|
|
|
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)
|
|
with_ignore_sid = parameterized.expand(
|
|
product(chain.from_iterable(asset_infos), [True, False])
|
|
)
|
|
|
|
|
|
def _utc_localize_index_level_0(df):
|
|
"""``tz_localize`` the first level of a multiindexed dataframe to utc.
|
|
|
|
Mutates df in place.
|
|
"""
|
|
idx = df.index
|
|
df.index = pd.MultiIndex.from_product(
|
|
(idx.levels[0].tz_localize('utc'), idx.levels[1]),
|
|
names=idx.names,
|
|
)
|
|
return df
|
|
|
|
|
|
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.),
|
|
'int_value': (0, 1, 2, 1, 2, 3, 2, 3, 4),
|
|
'asof_date': dates,
|
|
'timestamp': dates,
|
|
})
|
|
cls.dshape = dshape("""
|
|
var * {
|
|
sid: ?int64,
|
|
value: ?float64,
|
|
int_value: ?int64,
|
|
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()
|
|
cls.missing_values = {'int_value': 0}
|
|
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
self.assertEqual(ds.__name__, name)
|
|
self.assertTrue(issubclass(ds, DataSet))
|
|
|
|
self.assertIs(ds.value.dtype, float64_dtype)
|
|
self.assertIs(ds.int_value.dtype, int64_dtype)
|
|
|
|
self.assertTrue(np.isnan(ds.value.missing_value))
|
|
self.assertEqual(ds.int_value.missing_value, 0)
|
|
|
|
invalid_type_fields = ('asof_date',)
|
|
|
|
for field in invalid_type_fields:
|
|
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,
|
|
missing_values=self.missing_values,
|
|
),
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
self.assertEqual(value.name, 'value')
|
|
self.assertIsInstance(value, BoundColumn)
|
|
self.assertIs(value.dtype, float64_dtype)
|
|
|
|
# test memoization
|
|
self.assertIs(
|
|
from_blaze(
|
|
expr.value,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
missing_values=self.missing_values,
|
|
),
|
|
value,
|
|
)
|
|
self.assertIs(
|
|
from_blaze(
|
|
expr,
|
|
loader=self.garbage_loader,
|
|
no_deltas_rule=no_deltas_rules.ignore,
|
|
missing_values=self.missing_values,
|
|
).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,
|
|
missing_values=self.missing_values,
|
|
),
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
|
|
with self.assertRaises(TypeError):
|
|
from_blaze(
|
|
expr.value + 1, # put an Add in the column
|
|
deltas=None,
|
|
loader=self.garbage_loader,
|
|
missing_values=self.missing_values,
|
|
)
|
|
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
|
|
with self.assertRaises(TypeError):
|
|
from_blaze(
|
|
expr.value + 1,
|
|
deltas=deltas,
|
|
loader=self.garbage_loader,
|
|
missing_values=self.missing_values,
|
|
)
|
|
|
|
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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
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,
|
|
_utc_localize_index_level_0(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,
|
|
missing_values=self.missing_values,
|
|
)
|
|
p = Pipeline()
|
|
p.add(ds.value.latest, 'value')
|
|
p.add(ds.int_value.latest, 'int_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]',
|
|
).dt.tz_localize('utc')
|
|
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, ('int_value', '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', 'int_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']
|
|
|
|
asset_info = asset_infos[0][0]
|
|
with tmp_asset_finder(equities=asset_info) as finder:
|
|
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)
|
|
self._test_id(
|
|
df,
|
|
var * Record(fields),
|
|
expected,
|
|
finder,
|
|
('value', 'other'),
|
|
)
|
|
|
|
def test_id_take_last_in_group(self):
|
|
T = pd.Timestamp
|
|
df = pd.DataFrame(
|
|
columns=['asof_date', 'timestamp', 'sid', 'other', 'value'],
|
|
data=[
|
|
[T('2014-01-01'), T('2014-01-01 00'), 65, 0, 0],
|
|
[T('2014-01-01'), T('2014-01-01 01'), 65, 1, np.nan],
|
|
[T('2014-01-01'), T('2014-01-01 00'), 66, np.nan, np.nan],
|
|
[T('2014-01-01'), T('2014-01-01 01'), 66, np.nan, 1],
|
|
[T('2014-01-01'), T('2014-01-01 00'), 67, 2, np.nan],
|
|
[T('2014-01-01'), T('2014-01-01 01'), 67, np.nan, np.nan],
|
|
[T('2014-01-02'), T('2014-01-02 00'), 65, np.nan, np.nan],
|
|
[T('2014-01-02'), T('2014-01-02 01'), 65, np.nan, 1],
|
|
[T('2014-01-02'), T('2014-01-02 00'), 66, np.nan, np.nan],
|
|
[T('2014-01-02'), T('2014-01-02 01'), 66, 2, np.nan],
|
|
[T('2014-01-02'), T('2014-01-02 00'), 67, 3, 3],
|
|
[T('2014-01-02'), T('2014-01-02 01'), 67, 3, 3],
|
|
[T('2014-01-03'), T('2014-01-03 00'), 65, 2, np.nan],
|
|
[T('2014-01-03'), T('2014-01-03 01'), 65, 2, np.nan],
|
|
[T('2014-01-03'), T('2014-01-03 00'), 66, 3, 3],
|
|
[T('2014-01-03'), T('2014-01-03 01'), 66, np.nan, np.nan],
|
|
[T('2014-01-03'), T('2014-01-03 00'), 67, np.nan, np.nan],
|
|
[T('2014-01-03'), T('2014-01-03 01'), 67, np.nan, 4],
|
|
],
|
|
)
|
|
fields = OrderedDict(self.dshape.measure.fields)
|
|
fields['other'] = fields['value']
|
|
|
|
with tmp_asset_finder() as finder:
|
|
expected = pd.DataFrame(
|
|
columns=['other', 'value'],
|
|
data=[
|
|
[1, 0], # 2014-01-01 Equity(65 [A])
|
|
[np.nan, 1], # Equity(66 [B])
|
|
[2, np.nan], # Equity(67 [C])
|
|
[1, 1], # 2014-01-02 Equity(65 [A])
|
|
[2, 1], # Equity(66 [B])
|
|
[3, 3], # Equity(67 [C])
|
|
[2, 1], # 2014-01-03 Equity(65 [A])
|
|
[3, 3], # Equity(66 [B])
|
|
[3, 3], # Equity(67 [C])
|
|
],
|
|
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_take_last_in_group_macro(self):
|
|
"""
|
|
output (expected):
|
|
|
|
other value
|
|
2014-01-01 Equity(65 [A]) NaN 1
|
|
Equity(66 [B]) NaN 1
|
|
Equity(67 [C]) NaN 1
|
|
2014-01-02 Equity(65 [A]) 1 2
|
|
Equity(66 [B]) 1 2
|
|
Equity(67 [C]) 1 2
|
|
2014-01-03 Equity(65 [A]) 2 2
|
|
Equity(66 [B]) 2 2
|
|
Equity(67 [C]) 2 2
|
|
"""
|
|
T = pd.Timestamp
|
|
df = pd.DataFrame(
|
|
columns=['asof_date', 'timestamp', 'other', 'value'],
|
|
data=[
|
|
[T('2014-01-01'), T('2014-01-01 00'), np.nan, 1],
|
|
[T('2014-01-01'), T('2014-01-01 01'), np.nan, np.nan],
|
|
[T('2014-01-02'), T('2014-01-02 00'), 1, np.nan],
|
|
[T('2014-01-02'), T('2014-01-02 01'), np.nan, 2],
|
|
[T('2014-01-03'), T('2014-01-03 00'), 2, np.nan],
|
|
[T('2014-01-03'), T('2014-01-03 01'), 3, 3],
|
|
],
|
|
)
|
|
fields = OrderedDict(self.macro_dshape.measure.fields)
|
|
fields['other'] = fields['value']
|
|
|
|
with tmp_asset_finder() as finder:
|
|
expected = pd.DataFrame(
|
|
columns=[
|
|
'other', 'value',
|
|
],
|
|
data=[
|
|
[np.nan, 1], # 2014-01-01 Equity(65 [A])
|
|
[np.nan, 1], # Equity(66 [B])
|
|
[np.nan, 1], # Equity(67 [C])
|
|
[1, 2], # 2014-01-02 Equity(65 [A])
|
|
[1, 2], # Equity(66 [B])
|
|
[1, 2], # Equity(67 [C])
|
|
[2, 2], # 2014-01-03 Equity(65 [A])
|
|
[2, 2], # Equity(66 [B])
|
|
[2, 2], # Equity(67 [C])
|
|
],
|
|
index=pd.MultiIndex.from_product(
|
|
(self.dates, finder.retrieve_all(self.sids)),
|
|
),
|
|
)
|
|
self._test_id(
|
|
df,
|
|
var * Record(fields),
|
|
expected,
|
|
finder,
|
|
('value', 'other'),
|
|
)
|
|
|
|
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_,
|
|
missing_values=self.missing_values,
|
|
)
|
|
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,
|
|
_utc_localize_index_level_0(expected_output),
|
|
check_dtype=False,
|
|
)
|
|
|
|
@with_ignore_sid
|
|
def test_deltas(self, asset_info, add_extra_sid):
|
|
df = self.df.copy()
|
|
if add_extra_sid:
|
|
extra_sid_df = pd.DataFrame({
|
|
'asof_date': self.dates,
|
|
'timestamp': self.dates,
|
|
'sid': (ord('E'),) * 3,
|
|
'value': (3., 4., 5.,),
|
|
'int_value': (3, 4, 5),
|
|
})
|
|
df = df.append(extra_sid_df, ignore_index=True)
|
|
expr = bz.data(df, name='expr', dshape=self.dshape)
|
|
deltas = bz.data(df, dshape=self.dshape)
|
|
deltas = bz.data(
|
|
odo(
|
|
bz.transform(
|
|
deltas,
|
|
value=deltas.value + 10,
|
|
timestamp=deltas.timestamp + timedelta(days=1),
|
|
),
|
|
pd.DataFrame,
|
|
),
|
|
name='delta',
|
|
dshape=self.dshape,
|
|
)
|
|
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,
|
|
)
|
|
|
|
@with_extra_sid
|
|
def test_deltas_only_one_delta_in_universe(self, asset_info):
|
|
expr = bz.data(self.df, name='expr', dshape=self.dshape)
|
|
deltas = pd.DataFrame({
|
|
'sid': [65, 66],
|
|
'asof_date': [self.dates[1], self.dates[0]],
|
|
'timestamp': [self.dates[2], self.dates[1]],
|
|
'value': [10, 11],
|
|
})
|
|
deltas = bz.data(deltas, name='deltas', dshape=self.dshape)
|
|
expected_views = keymap(pd.Timestamp, {
|
|
'2014-01-02': np.array([[0.0, 11.0, 2.0],
|
|
[1.0, 2.0, 3.0]]),
|
|
'2014-01-03': np.array([[10.0, 2.0, 3.0],
|
|
[2.0, 3.0, 4.0]]),
|
|
'2014-01-04': np.array([[2.0, 3.0, 4.0],
|
|
[2.0, 3.0, 4.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(
|
|
columns=[
|
|
'value',
|
|
],
|
|
data=np.array([11, 10, 4]).repeat(len(asset_info.index)),
|
|
index=pd.MultiIndex.from_product((
|
|
sorted(expected_views.keys()),
|
|
finder.retrieve_all(asset_info.index),
|
|
)),
|
|
)
|
|
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.),
|
|
'int_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(
|
|
odo(
|
|
bz.transform(
|
|
expr,
|
|
value=expr.value + 10,
|
|
timestamp=expr.timestamp + timedelta(days=1),
|
|
),
|
|
pd.DataFrame,
|
|
),
|
|
name='delta',
|
|
dshape=self.dshape,
|
|
)
|
|
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),
|
|
)
|