mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-21 12:30:16 +08:00
ENH/BUG: Modeling API enhancements.
- Fixes an error where Modeling API data known as of the close of `day N` would be shown to algorithms during `before_trading_start` as of the close of the same day. Algorithms should now only receive data during `before_trading_start/handle_data` that was known as of the simulation time at which the function would be called. - All Term instances now have a `mask` attribute that must be a `Filter` or an instance of `AssetExists()`. `mask` can be used to specify that a Factor should be computed in a manner that ignores the values that were not `True` in the mask. - Changed the interface for `FFCLoader.load_adjusted_array` and `Term._compute` from `(columns, mask)`, with mask as a DataFrame, to `(columns, dates, assets, mask)`, where mask is a numpy array. This is primarily to avoid having to reconstruct extra DataFrames when using masks produced by non `AssetExists` filters. - Adds `BoundColumn.latest`, which gives the most-recently-known value of a column.
This commit is contained in:
@@ -64,8 +64,12 @@ ext_modules = LazyCythonizingList([
|
||||
('zipline.lib.adjustment', ['zipline/lib/adjustment.pyx']),
|
||||
('zipline.lib.rank', ['zipline/lib/rank.pyx']),
|
||||
(
|
||||
'zipline.data.ffc.loaders._us_equity_pricing',
|
||||
['zipline/data/ffc/loaders/_us_equity_pricing.pyx']
|
||||
'zipline.data.ffc.loaders._equities',
|
||||
['zipline/data/ffc/loaders/_equities.pyx'],
|
||||
),
|
||||
(
|
||||
'zipline.data.ffc.loaders._adjustments',
|
||||
['zipline/data/ffc/loaders/_adjustments.pyx'],
|
||||
),
|
||||
])
|
||||
|
||||
|
||||
+31
-13
@@ -5,13 +5,14 @@ from functools import wraps
|
||||
from unittest import TestCase
|
||||
|
||||
from numpy import arange, prod
|
||||
from numpy.random import randn, seed as random_seed
|
||||
from pandas import date_range, Int64Index, DataFrame
|
||||
from six import iteritems
|
||||
|
||||
from zipline.finance.trading import TradingEnvironment
|
||||
from zipline.modelling.engine import SimpleFFCEngine
|
||||
from zipline.modelling.graph import TermGraph
|
||||
from zipline.modelling.term import AssetExists
|
||||
from zipline.utils.pandas_utils import explode
|
||||
from zipline.utils.test_utils import make_simple_asset_info, ExplodingObject
|
||||
from zipline.utils.tradingcalendar import trading_day
|
||||
|
||||
@@ -56,9 +57,15 @@ class BaseFFCTestCase(TestCase):
|
||||
assets,
|
||||
self.__calendar[0],
|
||||
self.__calendar[-1],
|
||||
))
|
||||
),
|
||||
)
|
||||
self.__finder = env.asset_finder
|
||||
self.__mask = self.__finder.lifetimes(self.__calendar[-10:])
|
||||
|
||||
# Use a 30-day period at the end of the year by default.
|
||||
self.__mask = self.__finder.lifetimes(
|
||||
self.__calendar[-30:],
|
||||
include_start_date=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def default_shape(self):
|
||||
@@ -74,6 +81,12 @@ class BaseFFCTestCase(TestCase):
|
||||
----------
|
||||
terms : dict
|
||||
Mapping from termname -> term object.
|
||||
initial_workspace : dict
|
||||
Initial workspace to forward to SimpleFFCEngine.compute_chunk.
|
||||
mask : DataFrame, optional
|
||||
This is a value to pass to `initial_workspace` as the mask from
|
||||
`AssetExists()`. Defaults to a frame of shape `self.default_shape`
|
||||
containing all True values.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -85,10 +98,23 @@ class BaseFFCTestCase(TestCase):
|
||||
self.__calendar,
|
||||
self.__finder,
|
||||
)
|
||||
mask = mask if mask is not None else self.__mask
|
||||
return engine.compute_chunk(TermGraph(terms), mask, initial_workspace)
|
||||
if mask is None:
|
||||
mask = self.__mask
|
||||
|
||||
dates, assets, mask_values = explode(mask)
|
||||
initial_workspace.setdefault(AssetExists(), mask_values)
|
||||
return engine.compute_chunk(
|
||||
TermGraph(terms),
|
||||
dates,
|
||||
assets,
|
||||
initial_workspace,
|
||||
)
|
||||
|
||||
def build_mask(self, array):
|
||||
"""
|
||||
Helper for constructing an AssetExists mask from a boolean-coercible
|
||||
array.
|
||||
"""
|
||||
ndates, nassets = array.shape
|
||||
return DataFrame(
|
||||
array,
|
||||
@@ -105,11 +131,3 @@ class BaseFFCTestCase(TestCase):
|
||||
Build a block of testing data from numpy.arange.
|
||||
"""
|
||||
return arange(prod(shape), dtype=dtype).reshape(shape)
|
||||
|
||||
@with_default_shape
|
||||
def randn_data(self, seed, shape):
|
||||
"""
|
||||
Build a block of testing data from numpy.random.randn.
|
||||
"""
|
||||
random_seed(seed)
|
||||
return randn(*shape)
|
||||
|
||||
+113
-75
@@ -7,8 +7,8 @@ from itertools import product
|
||||
|
||||
from numpy import (
|
||||
full,
|
||||
isnan,
|
||||
nan,
|
||||
zeros,
|
||||
)
|
||||
from numpy.testing import assert_array_equal
|
||||
from pandas import (
|
||||
@@ -191,33 +191,21 @@ class ConstantInputTestCase(TestCase):
|
||||
expected_high_low = 3.0 * (constants[high] - constants[low])
|
||||
assert_frame_equal(
|
||||
high_low_result,
|
||||
DataFrame(
|
||||
expected_high_low,
|
||||
index=dates,
|
||||
columns=self.assets,
|
||||
)
|
||||
DataFrame(expected_high_low, index=dates, columns=self.assets),
|
||||
)
|
||||
|
||||
open_close_result = results['open_close'].unstack()
|
||||
expected_open_close = 3.0 * (constants[open] - constants[close])
|
||||
assert_frame_equal(
|
||||
open_close_result,
|
||||
DataFrame(
|
||||
expected_open_close,
|
||||
index=dates,
|
||||
columns=self.assets,
|
||||
)
|
||||
DataFrame(expected_open_close, index=dates, columns=self.assets),
|
||||
)
|
||||
|
||||
avg_result = results['avg'].unstack()
|
||||
expected_avg = (expected_high_low + expected_open_close) / 2.0
|
||||
assert_frame_equal(
|
||||
avg_result,
|
||||
DataFrame(
|
||||
expected_avg,
|
||||
index=dates,
|
||||
columns=self.assets,
|
||||
)
|
||||
DataFrame(expected_avg, index=dates, columns=self.assets),
|
||||
)
|
||||
|
||||
|
||||
@@ -338,20 +326,18 @@ class SyntheticBcolzTestCase(TestCase):
|
||||
def setUpClass(cls):
|
||||
cls.first_asset_start = Timestamp('2015-04-01', tz='UTC')
|
||||
cls.env = TradingEnvironment()
|
||||
cls.trading_day = cls.env.trading_day
|
||||
cls.trading_day = day = cls.env.trading_day
|
||||
cls.calendar = date_range('2015', '2015-08', tz='UTC', freq=day)
|
||||
|
||||
cls.asset_info = make_rotating_asset_info(
|
||||
num_assets=6,
|
||||
first_start=cls.first_asset_start,
|
||||
frequency=cls.trading_day,
|
||||
frequency=day,
|
||||
periods_between_starts=4,
|
||||
asset_lifetime=8,
|
||||
)
|
||||
cls.last_asset_end = cls.asset_info['end_date'].max()
|
||||
cls.all_assets = cls.asset_info.index
|
||||
cls.all_dates = date_range(
|
||||
start=cls.first_asset_start,
|
||||
end=cls.asset_info['end_date'].max(),
|
||||
freq=cls.trading_day,
|
||||
)
|
||||
|
||||
cls.env.write_data(equities_df=cls.asset_info)
|
||||
cls.finder = cls.env.asset_finder
|
||||
@@ -359,34 +345,74 @@ class SyntheticBcolzTestCase(TestCase):
|
||||
cls.temp_dir = TempDirectory()
|
||||
cls.temp_dir.create()
|
||||
|
||||
cls.writer = SyntheticDailyBarWriter(
|
||||
asset_info=cls.asset_info[['start_date', 'end_date']],
|
||||
calendar=cls.all_dates,
|
||||
)
|
||||
table = cls.writer.write(
|
||||
cls.temp_dir.getpath('testdata.bcolz'),
|
||||
cls.all_dates,
|
||||
cls.all_assets,
|
||||
)
|
||||
try:
|
||||
cls.writer = SyntheticDailyBarWriter(
|
||||
asset_info=cls.asset_info[['start_date', 'end_date']],
|
||||
calendar=cls.calendar,
|
||||
)
|
||||
table = cls.writer.write(
|
||||
cls.temp_dir.getpath('testdata.bcolz'),
|
||||
cls.calendar,
|
||||
cls.all_assets,
|
||||
)
|
||||
|
||||
cls.ffc_loader = USEquityPricingLoader(
|
||||
BcolzDailyBarReader(table),
|
||||
NullAdjustmentReader(),
|
||||
)
|
||||
cls.ffc_loader = USEquityPricingLoader(
|
||||
BcolzDailyBarReader(table),
|
||||
NullAdjustmentReader(),
|
||||
)
|
||||
except:
|
||||
cls.temp_dir.cleanup()
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
del cls.env
|
||||
cls.temp_dir.cleanup()
|
||||
|
||||
def write_nans(self, df):
|
||||
"""
|
||||
Write nans to the locations in data corresponding to the (date, asset)
|
||||
pairs for which we wouldn't have data for `asset` on `date` in a
|
||||
backtest.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
df : pd.DataFrame
|
||||
A DataFrame with a DatetimeIndex as index and an object index of
|
||||
Assets as columns.
|
||||
|
||||
This means that we write nans for dates after an asset's end_date and
|
||||
**on or before** an asset's start_date. The assymetry here is because
|
||||
of the fact that, on the morning of an asset's first date, we haven't
|
||||
yet seen any trades for that asset, so we wouldn't be able to show any
|
||||
useful data to the user.
|
||||
"""
|
||||
# Mask out with nans all the dates on which each asset didn't exist
|
||||
index = df.index
|
||||
min_, max_ = index[[0, -1]]
|
||||
for asset in df.columns:
|
||||
if asset.start_date >= min_:
|
||||
start = index.get_loc(asset.start_date, method='bfill')
|
||||
df.loc[:start + 1, asset] = nan # +1 to overwrite start_date
|
||||
if asset.end_date <= max_:
|
||||
end = index.get_loc(asset.end_date)
|
||||
df.ix[end + 1:, asset] = nan # +1 to *not* overwrite end_date
|
||||
|
||||
def test_SMA(self):
|
||||
engine = SimpleFFCEngine(
|
||||
self.ffc_loader,
|
||||
self.env.trading_days,
|
||||
self.finder,
|
||||
)
|
||||
dates, assets = self.all_dates, self.all_assets
|
||||
window_length = 5
|
||||
assets = self.all_assets
|
||||
dates = date_range(
|
||||
self.first_asset_start + self.trading_day,
|
||||
self.last_asset_end,
|
||||
freq=self.trading_day,
|
||||
)
|
||||
dates_to_test = dates[window_length:]
|
||||
|
||||
SMA = SimpleMovingAverage(
|
||||
inputs=(USEquityPricing.close,),
|
||||
window_length=window_length,
|
||||
@@ -394,27 +420,30 @@ class SyntheticBcolzTestCase(TestCase):
|
||||
|
||||
results = engine.factor_matrix(
|
||||
{'sma': SMA},
|
||||
dates[window_length],
|
||||
dates[-1],
|
||||
dates_to_test[0],
|
||||
dates_to_test[-1],
|
||||
)
|
||||
raw_closes = self.writer.expected_values_2d(dates, assets, 'close')
|
||||
expected_sma_result = rolling_mean(
|
||||
raw_closes,
|
||||
|
||||
# Shift back the raw inputs by a trading day because we expect our
|
||||
# computed results to be computed using values anchored on the
|
||||
# **previous** day's data.
|
||||
expected_raw = rolling_mean(
|
||||
self.writer.expected_values_2d(
|
||||
dates - self.trading_day, assets, 'close',
|
||||
),
|
||||
window_length,
|
||||
min_periods=1,
|
||||
)
|
||||
expected_sma_result[isnan(raw_closes)] = nan
|
||||
expected_sma_result = expected_sma_result[window_length:]
|
||||
|
||||
sma_result = results['sma'].unstack()
|
||||
assert_frame_equal(
|
||||
sma_result,
|
||||
DataFrame(
|
||||
expected_sma_result,
|
||||
index=dates[window_length:],
|
||||
columns=assets,
|
||||
),
|
||||
expected = DataFrame(
|
||||
# Truncate off the extra rows needed to compute the SMAs.
|
||||
expected_raw[window_length:],
|
||||
index=dates_to_test, # dates_to_test is dates[window_length:]
|
||||
columns=self.finder.retrieve_all(assets),
|
||||
)
|
||||
self.write_nans(expected)
|
||||
result = results['sma'].unstack()
|
||||
assert_frame_equal(result, expected)
|
||||
|
||||
def test_drawdown(self):
|
||||
# The monotonically-increasing data produced by SyntheticDailyBarWriter
|
||||
@@ -427,8 +456,15 @@ class SyntheticBcolzTestCase(TestCase):
|
||||
self.env.trading_days,
|
||||
self.finder,
|
||||
)
|
||||
dates, assets = self.all_dates, self.all_assets
|
||||
window_length = 5
|
||||
assets = self.all_assets
|
||||
dates = date_range(
|
||||
self.first_asset_start + self.trading_day,
|
||||
self.last_asset_end,
|
||||
freq=self.trading_day,
|
||||
)
|
||||
dates_to_test = dates[window_length:]
|
||||
|
||||
drawdown = MaxDrawdown(
|
||||
inputs=(USEquityPricing.close,),
|
||||
window_length=window_length,
|
||||
@@ -436,32 +472,27 @@ class SyntheticBcolzTestCase(TestCase):
|
||||
|
||||
results = engine.factor_matrix(
|
||||
{'drawdown': drawdown},
|
||||
dates[window_length],
|
||||
dates[-1],
|
||||
dates_to_test[0],
|
||||
dates_to_test[-1],
|
||||
)
|
||||
|
||||
dd_result = results['drawdown']
|
||||
|
||||
# We expect NaNs when the asset was undefined, otherwise 0 everywhere,
|
||||
# since the input is always increasing.
|
||||
expected = self.writer.expected_values_2d(dates, assets, 'close')
|
||||
expected[~isnan(expected)] = 0
|
||||
expected = expected[window_length:]
|
||||
|
||||
assert_frame_equal(
|
||||
dd_result.unstack(),
|
||||
DataFrame(
|
||||
expected,
|
||||
index=dates[window_length:],
|
||||
columns=assets,
|
||||
),
|
||||
expected = DataFrame(
|
||||
data=zeros((len(dates_to_test), len(assets)), dtype=float),
|
||||
index=dates_to_test,
|
||||
columns=self.finder.retrieve_all(assets),
|
||||
)
|
||||
self.write_nans(expected)
|
||||
result = results['drawdown'].unstack()
|
||||
|
||||
assert_frame_equal(expected, result)
|
||||
|
||||
|
||||
class MultiColumnLoaderTestCase(TestCase):
|
||||
def setUp(self):
|
||||
self.assets = [1, 2, 3]
|
||||
self.dates = date_range('2014-01-01', '2014-02-01', freq='D', tz='UTC')
|
||||
self.dates = date_range('2014-01', '2014-03', freq='D', tz='UTC')
|
||||
|
||||
asset_info = make_simple_asset_info(
|
||||
self.assets,
|
||||
@@ -475,6 +506,11 @@ class MultiColumnLoaderTestCase(TestCase):
|
||||
def test_engine_with_multicolumn_loader(self):
|
||||
open_, close = USEquityPricing.open, USEquityPricing.close
|
||||
|
||||
# Test for thirty days up to the second to last day that we think all
|
||||
# the assets existed. If we test the last day of our calendar, no
|
||||
# assets will be in our output, because their end dates are all
|
||||
dates_to_test = self.dates[-32:-2]
|
||||
|
||||
loader = MultiColumnLoader({
|
||||
open_: ConstantLoader(dates=self.dates,
|
||||
assets=self.assets,
|
||||
@@ -489,12 +525,14 @@ class MultiColumnLoaderTestCase(TestCase):
|
||||
factor = RollingSumDifference()
|
||||
|
||||
result = engine.factor_matrix({'f': factor},
|
||||
self.dates[2],
|
||||
self.dates[-1])
|
||||
dates_to_test[0],
|
||||
dates_to_test[-1])
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual({'f'}, set(result.columns))
|
||||
|
||||
# (close - open) * window = (1 - 2) * 3 = -3
|
||||
# skipped 2 from the start, so that the window is full
|
||||
check_arrays(result['f'],
|
||||
Series([-3] * len(self.assets) * (len(self.dates) - 2)))
|
||||
result_index = self.assets * len(dates_to_test)
|
||||
result_shape = (len(result_index),)
|
||||
check_arrays(
|
||||
result['f'],
|
||||
Series(index=result_index, data=full(result_shape, -3)),
|
||||
)
|
||||
|
||||
@@ -1,13 +1,10 @@
|
||||
"""
|
||||
Tests for Factor terms.
|
||||
"""
|
||||
from numpy import (
|
||||
array,
|
||||
ones,
|
||||
)
|
||||
|
||||
from numpy import array, eye, nan, ones
|
||||
from zipline.errors import UnknownRankMethod
|
||||
from zipline.modelling.factor import Factor
|
||||
from zipline.modelling.filter import Filter
|
||||
from zipline.utils.test_utils import check_arrays
|
||||
|
||||
from .base import BaseFFCTestCase
|
||||
@@ -18,6 +15,11 @@ class F(Factor):
|
||||
window_length = 0
|
||||
|
||||
|
||||
class Mask(Filter):
|
||||
inputs = ()
|
||||
window_length = 0
|
||||
|
||||
|
||||
class FactorTestCase(BaseFFCTestCase):
|
||||
|
||||
def setUp(self):
|
||||
@@ -135,3 +137,43 @@ class FactorTestCase(BaseFFCTestCase):
|
||||
})
|
||||
# Not passing a method should default to ordinal.
|
||||
check({'ordinal': self.f.rank(ascending=False)})
|
||||
|
||||
def test_rank_after_mask(self):
|
||||
# data = arange(25).reshape(5, 5).transpose() % 4
|
||||
data = array([[0, 1, 2, 3, 0],
|
||||
[1, 2, 3, 0, 1],
|
||||
[2, 3, 0, 1, 2],
|
||||
[3, 0, 1, 2, 3],
|
||||
[0, 1, 2, 3, 0]], dtype=float)
|
||||
mask_data = ~eye(5, dtype=bool)
|
||||
initial_workspace = {self.f: data, Mask(): mask_data}
|
||||
|
||||
terms = {
|
||||
"ascending_nomask": self.f.rank(ascending=True),
|
||||
"ascending_mask": self.f.rank(ascending=True, mask=Mask()),
|
||||
# "descending_nomask": self.f.rank(ascending=False),
|
||||
# "descending_mask": self.f.rank(ascending=False, mask=Mask()),
|
||||
}
|
||||
|
||||
expected = {
|
||||
"ascending_nomask": array([[1., 3., 4., 5., 2.],
|
||||
[2., 4., 5., 1., 3.],
|
||||
[3., 5., 1., 2., 4.],
|
||||
[4., 1., 2., 3., 5.],
|
||||
[1., 3., 4., 5., 2.]]),
|
||||
# Diagonal should be all nans, and anything whose rank was less
|
||||
# than the diagonal in the unmasked calc should go down by 1.
|
||||
"ascending_mask": array([[nan, 2., 3., 4., 1.],
|
||||
[2., nan, 4., 1., 3.],
|
||||
[2., 4., nan, 1., 3.],
|
||||
[3., 1., 2., nan, 4.],
|
||||
[1., 2., 3., 4., nan]]),
|
||||
}
|
||||
|
||||
results = self.run_terms(
|
||||
terms,
|
||||
initial_workspace,
|
||||
mask=self.build_mask(ones((5, 5))),
|
||||
)
|
||||
for method in results:
|
||||
check_arrays(expected[method], results[method])
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""
|
||||
Tests for filter terms.
|
||||
"""
|
||||
from itertools import product
|
||||
from operator import and_
|
||||
|
||||
from numpy import (
|
||||
@@ -13,17 +14,20 @@ from numpy import (
|
||||
nan,
|
||||
nanpercentile,
|
||||
ones,
|
||||
ones_like,
|
||||
putmask,
|
||||
)
|
||||
from numpy.random import randn, seed as random_seed
|
||||
|
||||
from zipline.errors import BadPercentileBounds
|
||||
from zipline.modelling.filter import Filter
|
||||
from zipline.modelling.factor import Factor
|
||||
from zipline.utils.test_utils import check_arrays
|
||||
|
||||
from .base import BaseFFCTestCase
|
||||
from .base import BaseFFCTestCase, with_default_shape
|
||||
|
||||
|
||||
def rowwise_rank(array):
|
||||
def rowwise_rank(array, mask=None):
|
||||
"""
|
||||
Take a 2D array and return the 0-indexed sorted position of each element in
|
||||
the array for each row.
|
||||
@@ -61,6 +65,11 @@ class SomeOtherFactor(Factor):
|
||||
window_length = 0
|
||||
|
||||
|
||||
class Mask(Filter):
|
||||
inputs = ()
|
||||
window_length = 0
|
||||
|
||||
|
||||
class FilterTestCase(BaseFFCTestCase):
|
||||
|
||||
def setUp(self):
|
||||
@@ -68,6 +77,14 @@ class FilterTestCase(BaseFFCTestCase):
|
||||
self.f = SomeFactor()
|
||||
self.g = SomeOtherFactor()
|
||||
|
||||
@with_default_shape
|
||||
def randn_data(self, seed, shape):
|
||||
"""
|
||||
Build a block of testing data from numpy.random.randn.
|
||||
"""
|
||||
random_seed(seed)
|
||||
return randn(*shape)
|
||||
|
||||
def test_bad_percentiles(self):
|
||||
f = self.f
|
||||
|
||||
@@ -81,25 +98,58 @@ class FilterTestCase(BaseFFCTestCase):
|
||||
with self.assertRaises(BadPercentileBounds):
|
||||
f.percentile_between(min_, max_)
|
||||
|
||||
def test_top(self):
|
||||
def test_top_and_bottom(self):
|
||||
data = self.randn_data(seed=5) # Fix a seed for determinism.
|
||||
|
||||
mask_data = ones_like(data, dtype=bool)
|
||||
mask_data[:, 0] = False
|
||||
|
||||
nan_data = data.copy()
|
||||
nan_data[:, 0] = nan
|
||||
|
||||
mask = Mask()
|
||||
initial_workspace = {self.f: data, mask: mask_data}
|
||||
|
||||
methods = ['top', 'bottom']
|
||||
counts = 2, 3, 10
|
||||
data = self.randn_data(seed=5) # Arbitrary seed choice.
|
||||
results = self.run_terms(
|
||||
terms={'top_' + str(c): self.f.top(c) for c in counts},
|
||||
initial_workspace={self.f: data},
|
||||
)
|
||||
for c in counts:
|
||||
result = results['top_' + str(c)]
|
||||
term_combos = list(product(methods, counts, [True, False]))
|
||||
|
||||
def termname(method, count, masked):
|
||||
return '_'.join([method, str(count), 'mask' if masked else ''])
|
||||
|
||||
# Add a term for each permutation of top/bottom, count, and
|
||||
# mask/no_mask.
|
||||
terms = {}
|
||||
for method, count, masked in term_combos:
|
||||
kwargs = {'N': count}
|
||||
if masked:
|
||||
kwargs['mask'] = mask
|
||||
term = getattr(self.f, method)(**kwargs)
|
||||
terms[termname(method, count, masked)] = term
|
||||
|
||||
results = self.run_terms(terms, initial_workspace=initial_workspace)
|
||||
|
||||
def expected_result(method, count, masked):
|
||||
# Ranking with a mask is equivalent to ranking with nans applied on
|
||||
# the masked values.
|
||||
to_rank = nan_data if masked else data
|
||||
|
||||
if method == 'top':
|
||||
return rowwise_rank(-to_rank) < count
|
||||
elif method == 'bottom':
|
||||
return rowwise_rank(to_rank) < count
|
||||
|
||||
for method, count, masked in term_combos:
|
||||
result = results[termname(method, count, masked)]
|
||||
|
||||
# Check that `min(c, num_assets)` assets passed each day.
|
||||
passed_per_day = result.sum(axis=1)
|
||||
check_arrays(
|
||||
passed_per_day,
|
||||
full_like(passed_per_day, min(c, data.shape[1])),
|
||||
full_like(passed_per_day, min(count, data.shape[1])),
|
||||
)
|
||||
|
||||
# Check that the top `c` assets passed.
|
||||
expected = rowwise_rank(-data) < c
|
||||
expected = expected_result(method, count, masked)
|
||||
check_arrays(result, expected)
|
||||
|
||||
def test_bottom(self):
|
||||
@@ -228,35 +278,19 @@ class FilterTestCase(BaseFFCTestCase):
|
||||
)
|
||||
check_arrays(result, expected)
|
||||
|
||||
def test_sequenced_filter_order_independent(self):
|
||||
data = self.arange_data() % 5
|
||||
results = self.run_terms(
|
||||
{
|
||||
# Sequencing is equivalent to &ing for commutative filters.
|
||||
'sequenced': (1.5 < self.f).then(self.f < 3.5),
|
||||
'anded': (1.5 < self.f) & (self.f < 3.5),
|
||||
},
|
||||
initial_workspace={self.f: data},
|
||||
)
|
||||
expected = (1.5 < data) & (data < 3.5)
|
||||
|
||||
check_arrays(results['sequenced'], expected)
|
||||
check_arrays(results['anded'], expected)
|
||||
|
||||
def test_sequenced_filter_order_dependent(self):
|
||||
|
||||
first = self.f < 1
|
||||
def test_percentile_after_mask(self):
|
||||
f_input = eye(5)
|
||||
|
||||
second = self.g.percentile_between(80, 100)
|
||||
g_input = arange(25, dtype=float).reshape(5, 5)
|
||||
|
||||
initial_mask = self.build_mask(ones((5, 5)))
|
||||
|
||||
custom_mask = self.f < 1
|
||||
without_mask = self.g.percentile_between(80, 100)
|
||||
with_mask = self.g.percentile_between(80, 100, mask=custom_mask)
|
||||
|
||||
terms = {
|
||||
'first': first,
|
||||
'second': second,
|
||||
'sequenced': first.then(second),
|
||||
'custom_mask': custom_mask,
|
||||
'without': without_mask,
|
||||
'with': with_mask,
|
||||
}
|
||||
|
||||
results = self.run_terms(
|
||||
@@ -266,11 +300,11 @@ class FilterTestCase(BaseFFCTestCase):
|
||||
)
|
||||
|
||||
# First should pass everything but the diagonal.
|
||||
check_arrays(results['first'], ~eye(5, dtype=bool))
|
||||
check_arrays(results['custom_mask'], ~eye(5, dtype=bool))
|
||||
|
||||
# Second should pass the largest value each day. Each row is strictly
|
||||
# increasing, so we always select the last value.
|
||||
expected_second = array(
|
||||
expected_without = array(
|
||||
[[0, 0, 0, 0, 1],
|
||||
[0, 0, 0, 0, 1],
|
||||
[0, 0, 0, 0, 1],
|
||||
@@ -278,12 +312,12 @@ class FilterTestCase(BaseFFCTestCase):
|
||||
[0, 0, 0, 0, 1]],
|
||||
dtype=bool,
|
||||
)
|
||||
check_arrays(results['second'], expected_second)
|
||||
check_arrays(results['without'], expected_without)
|
||||
|
||||
# When sequencing, we should remove the diagonal as an option before
|
||||
# computing percentiles. On the last day, we should get the
|
||||
# second-largest value, rather than the largest.
|
||||
expected_sequenced = array(
|
||||
expected_with = array(
|
||||
[[0, 0, 0, 0, 1],
|
||||
[0, 0, 0, 0, 1],
|
||||
[0, 0, 0, 0, 1],
|
||||
@@ -291,4 +325,4 @@ class FilterTestCase(BaseFFCTestCase):
|
||||
[0, 0, 0, 1, 0]], # Different from previous!
|
||||
dtype=bool,
|
||||
)
|
||||
check_arrays(results['sequenced'], expected_sequenced)
|
||||
check_arrays(results['with'], expected_with)
|
||||
|
||||
@@ -4,7 +4,7 @@ Tests for zipline.data.ffc.frame.DataFrameFFCLoader
|
||||
from unittest import TestCase
|
||||
|
||||
from mock import patch
|
||||
from numpy import arange
|
||||
from numpy import arange, ones
|
||||
from numpy.testing import assert_array_equal
|
||||
from pandas import (
|
||||
DataFrame,
|
||||
@@ -40,12 +40,7 @@ class DataFrameFFCLoaderTestCase(TestCase):
|
||||
periods=self.ndates,
|
||||
)
|
||||
|
||||
self.mask = DataFrame(
|
||||
True,
|
||||
index=self.dates,
|
||||
columns=self.sids,
|
||||
dtype=bool,
|
||||
)
|
||||
self.mask = ones((len(self.dates), len(self.sids)), dtype=bool)
|
||||
|
||||
def tearDown(self):
|
||||
pass
|
||||
@@ -60,13 +55,17 @@ class DataFrameFFCLoaderTestCase(TestCase):
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
# Wrong column.
|
||||
loader.load_adjusted_array([USEquityPricing.open], self.mask)
|
||||
loader.load_adjusted_array(
|
||||
[USEquityPricing.open], self.dates, self.sids, self.mask
|
||||
)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
# Too many columns.
|
||||
loader.load_adjusted_array(
|
||||
[USEquityPricing.open, USEquityPricing.close],
|
||||
self.mask
|
||||
self.dates,
|
||||
self.sids,
|
||||
self.mask,
|
||||
)
|
||||
|
||||
def test_baseline(self):
|
||||
@@ -81,7 +80,9 @@ class DataFrameFFCLoaderTestCase(TestCase):
|
||||
sids_slice = slice(1, 3, None)
|
||||
[adj_array] = loader.load_adjusted_array(
|
||||
[USEquityPricing.close],
|
||||
self.mask.iloc[dates_slice, sids_slice]
|
||||
self.dates[dates_slice],
|
||||
self.sids[sids_slice],
|
||||
self.mask[dates_slice, sids_slice],
|
||||
)
|
||||
|
||||
for idx, window in enumerate(adj_array.traverse(window_length=3)):
|
||||
@@ -203,10 +204,12 @@ class DataFrameFFCLoaderTestCase(TestCase):
|
||||
}
|
||||
self.assertEqual(formatted_adjustments, expected_formatted_adjustments)
|
||||
|
||||
mask = self.mask.iloc[dates_slice, sids_slice]
|
||||
mask = self.mask[dates_slice, sids_slice]
|
||||
with patch('zipline.data.ffc.frame.adjusted_array') as m:
|
||||
loader.load_adjusted_array(
|
||||
columns=[USEquityPricing.close],
|
||||
dates=self.dates[dates_slice],
|
||||
assets=self.sids[sids_slice],
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
@@ -214,5 +217,5 @@ class DataFrameFFCLoaderTestCase(TestCase):
|
||||
|
||||
args, kwargs = m.call_args
|
||||
assert_array_equal(kwargs['data'], expected_baseline.values)
|
||||
assert_array_equal(kwargs['mask'], mask.values)
|
||||
assert_array_equal(kwargs['mask'], mask)
|
||||
self.assertEqual(kwargs['adjustments'], expected_formatted_adjustments)
|
||||
|
||||
@@ -7,9 +7,9 @@ from os.path import (
|
||||
join,
|
||||
realpath,
|
||||
)
|
||||
|
||||
from numpy import (
|
||||
array,
|
||||
arange,
|
||||
full_like,
|
||||
nan,
|
||||
)
|
||||
@@ -17,22 +17,23 @@ from numpy.testing import assert_almost_equal
|
||||
from pandas import (
|
||||
concat,
|
||||
DataFrame,
|
||||
date_range,
|
||||
DatetimeIndex,
|
||||
Panel,
|
||||
read_csv,
|
||||
Series,
|
||||
Timestamp,
|
||||
)
|
||||
from six import iteritems
|
||||
from six import iteritems, itervalues
|
||||
from testfixtures import TempDirectory
|
||||
|
||||
from zipline.algorithm import TradingAlgorithm
|
||||
from zipline.api import (
|
||||
# add_filter,
|
||||
add_factor,
|
||||
get_datetime,
|
||||
)
|
||||
# from zipline.data.equities import USEquityPricing
|
||||
from zipline.data.equities import USEquityPricing
|
||||
from zipline.data.ffc.frame import DataFrameFFCLoader, MULTIPLY
|
||||
from zipline.data.ffc.loaders.us_equity_pricing import (
|
||||
BcolzDailyBarReader,
|
||||
DailyBarWriterFromCSVs,
|
||||
@@ -40,14 +41,17 @@ from zipline.data.ffc.loaders.us_equity_pricing import (
|
||||
SQLiteAdjustmentWriter,
|
||||
USEquityPricingLoader,
|
||||
)
|
||||
# from zipline.modelling.factor import CustomFactor
|
||||
from zipline.finance import trading
|
||||
|
||||
from zipline.modelling.factor.technical import VWAP
|
||||
from zipline.utils.test_utils import (
|
||||
make_simple_asset_info,
|
||||
str_to_seconds,
|
||||
)
|
||||
from zipline.utils.tradingcalendar import trading_days
|
||||
from zipline.utils.tradingcalendar import (
|
||||
trading_day,
|
||||
trading_days,
|
||||
)
|
||||
|
||||
|
||||
TEST_RESOURCE_PATH = join(
|
||||
@@ -70,6 +74,127 @@ def rolling_vwap(df, length):
|
||||
return Series(out, index=df.index)
|
||||
|
||||
|
||||
class ClosesOnly(TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.env = env = trading.TradingEnvironment()
|
||||
self.dates = date_range(
|
||||
'2014-01-01', '2014-02-01', freq=trading_day, tz='UTC'
|
||||
)
|
||||
asset_info = DataFrame.from_records([
|
||||
{
|
||||
'sid': 1,
|
||||
'symbol': 'A',
|
||||
'asset_type': 'equity',
|
||||
'start_date': self.dates[10],
|
||||
'end_date': self.dates[13],
|
||||
'exchange': 'TEST',
|
||||
},
|
||||
{
|
||||
'sid': 2,
|
||||
'symbol': 'B',
|
||||
'asset_type': 'equity',
|
||||
'start_date': self.dates[11],
|
||||
'end_date': self.dates[14],
|
||||
'exchange': 'TEST',
|
||||
},
|
||||
{
|
||||
'sid': 3,
|
||||
'symbol': 'C',
|
||||
'asset_type': 'equity',
|
||||
'start_date': self.dates[12],
|
||||
'end_date': self.dates[15],
|
||||
'exchange': 'TEST',
|
||||
},
|
||||
])
|
||||
self.first_asset_start = min(asset_info.start_date)
|
||||
self.last_asset_end = max(asset_info.end_date)
|
||||
env.write_data(equities_df=asset_info)
|
||||
self.asset_finder = finder = env.asset_finder
|
||||
|
||||
sids = (1, 2, 3)
|
||||
self.assets = finder.retrieve_all(sids)
|
||||
|
||||
# View of the baseline data.
|
||||
self.closes = DataFrame(
|
||||
{sid: arange(1, len(self.dates) + 1) * sid for sid in sids},
|
||||
index=self.dates,
|
||||
dtype=float,
|
||||
)
|
||||
|
||||
# Add a split for 'A' on its second date.
|
||||
self.split_asset = self.assets[0]
|
||||
self.split_date = self.split_asset.start_date + trading_day
|
||||
self.split_ratio = 0.5
|
||||
self.adjustments = DataFrame.from_records([
|
||||
{
|
||||
'sid': self.split_asset.sid,
|
||||
'value': self.split_ratio,
|
||||
'kind': MULTIPLY,
|
||||
'start_date': Timestamp('NaT'),
|
||||
'end_date': self.split_date,
|
||||
'apply_date': self.split_date,
|
||||
}
|
||||
])
|
||||
|
||||
# View of the data on/after the split.
|
||||
self.adj_closes = adj_closes = self.closes.copy()
|
||||
adj_closes.ix[:self.split_date, self.split_asset] *= self.split_ratio
|
||||
|
||||
self.ffc_loader = DataFrameFFCLoader(
|
||||
column=USEquityPricing.close,
|
||||
baseline=self.closes,
|
||||
adjustments=self.adjustments,
|
||||
)
|
||||
|
||||
def expected_close(self, date, asset):
|
||||
if date < self.split_date:
|
||||
lookup = self.closes
|
||||
else:
|
||||
lookup = self.adj_closes
|
||||
return lookup.loc[date, asset]
|
||||
|
||||
def exists(self, date, asset):
|
||||
return asset.start_date <= date <= asset.end_date
|
||||
|
||||
def test_assets_appear_on_correct_days(self):
|
||||
"""
|
||||
Assert that assets appear at correct times during a backtest, with
|
||||
correctly-adjusted close price values.
|
||||
"""
|
||||
def initialize(context):
|
||||
add_factor(USEquityPricing.close.latest, 'close')
|
||||
|
||||
def handle_data(context, data):
|
||||
factors = data.factors
|
||||
date = get_datetime().normalize()
|
||||
for asset in self.assets:
|
||||
# Assets should appear iff they exist today and yesterday.
|
||||
exists_today = self.exists(date, asset)
|
||||
existed_yesterday = self.exists(date - trading_day, asset)
|
||||
if exists_today and existed_yesterday:
|
||||
latest = factors.loc[asset, 'close']
|
||||
self.assertEqual(latest, self.expected_close(date, asset))
|
||||
else:
|
||||
self.assertNotIn(asset, factors.index)
|
||||
|
||||
before_trading_start = handle_data
|
||||
|
||||
algo = TradingAlgorithm(
|
||||
initialize=initialize,
|
||||
handle_data=handle_data,
|
||||
before_trading_start=before_trading_start,
|
||||
data_frequency='daily',
|
||||
ffc_loader=self.ffc_loader,
|
||||
start=self.first_asset_start - trading_day,
|
||||
end=self.last_asset_end + trading_day,
|
||||
env=self.env,
|
||||
)
|
||||
|
||||
# Run for a week in the middle of our data.
|
||||
algo.run(source=self.closes.iloc[10:17])
|
||||
|
||||
|
||||
class FFCAlgorithmTestCase(TestCase):
|
||||
|
||||
@classmethod
|
||||
@@ -155,16 +280,21 @@ class FFCAlgorithmTestCase(TestCase):
|
||||
|
||||
def test_handle_adjustment(self):
|
||||
AAPL, MSFT, BRK_A = assets = self.AAPL, self.MSFT, self.BRK_A
|
||||
raw_data = self.raw_data
|
||||
adjusted_data = {k: v.copy() for k, v in iteritems(raw_data)}
|
||||
|
||||
AAPL_split_date = Timestamp("2014-06-09", tz='UTC')
|
||||
split_loc = raw_data[AAPL].index.get_loc(AAPL_split_date)
|
||||
# Our view of the data before AAPL's split on June 9, 2014.
|
||||
raw = {k: v.copy() for k, v in iteritems(self.raw_data)}
|
||||
|
||||
# Our view of AAPL's history changes after the split.
|
||||
ohlc = ['open', 'high', 'low', 'close']
|
||||
adjusted_data[AAPL].ix[:split_loc, ohlc] /= 7.0
|
||||
adjusted_data[AAPL].ix[:split_loc, ['volume']] *= 7.0
|
||||
split_date = Timestamp("2014-06-09", tz='UTC')
|
||||
split_loc = self.dates.get_loc(split_date)
|
||||
split_ratio = 7.0
|
||||
|
||||
# Our view of the data after AAPL's split. All prices from before June
|
||||
# 9 get divided by the split ratio, and volumes get multiplied by the
|
||||
# split ratio.
|
||||
adj = {k: v.copy() for k, v in iteritems(self.raw_data)}
|
||||
for column in 'open', 'high', 'low', 'close':
|
||||
adj[AAPL].ix[:split_loc, column] /= split_ratio
|
||||
adj[AAPL].ix[:split_loc, 'volume'] *= split_ratio
|
||||
|
||||
window_lengths = [1, 2, 5, 10]
|
||||
# length -> asset -> expected vwap
|
||||
@@ -173,14 +303,56 @@ class FFCAlgorithmTestCase(TestCase):
|
||||
for length in window_lengths:
|
||||
vwap_keys[length] = "vwap_%d" % length
|
||||
for asset in AAPL, MSFT, BRK_A:
|
||||
raw = rolling_vwap(raw_data[asset], length)
|
||||
adj = rolling_vwap(adjusted_data[asset], length)
|
||||
raw_vwap = rolling_vwap(raw[asset], length)
|
||||
adj_vwap = rolling_vwap(adj[asset], length)
|
||||
# Shift computed results one day forward so that they're
|
||||
# labelled by the date on which they'll be seen in the
|
||||
# algorithm. (We can't show the close price for day N until day
|
||||
# N + 1.)
|
||||
vwaps[length][asset] = concat(
|
||||
[
|
||||
raw[:split_loc],
|
||||
adj[split_loc:]
|
||||
raw_vwap[:split_loc - 1],
|
||||
adj_vwap[split_loc - 1:]
|
||||
]
|
||||
)
|
||||
).shift(1, trading_day)
|
||||
|
||||
# Make sure all the expected vwaps have the same dates.
|
||||
vwap_dates = vwaps[1][self.AAPL].index
|
||||
for dict_ in itervalues(vwaps):
|
||||
# Each value is a dict mapping sid -> expected series.
|
||||
for series in itervalues(dict_):
|
||||
self.assertTrue((vwap_dates == series.index).all())
|
||||
|
||||
# Spot check expectations near the AAPL split.
|
||||
# length 1 vwap for the morning before the split should be the close
|
||||
# price of the previous day.
|
||||
before_split = vwaps[1][AAPL].loc[split_date - trading_day]
|
||||
assert_almost_equal(before_split, 647.3499, decimal=2)
|
||||
assert_almost_equal(
|
||||
before_split,
|
||||
raw[AAPL].loc[split_date - (2 * trading_day), 'close'],
|
||||
decimal=2,
|
||||
)
|
||||
|
||||
# length 1 vwap for the morning of the split should be the close price
|
||||
# of the previous day, **ADJUSTED FOR THE SPLIT**.
|
||||
on_split = vwaps[1][AAPL].loc[split_date]
|
||||
assert_almost_equal(on_split, 645.5700 / split_ratio, decimal=2)
|
||||
assert_almost_equal(
|
||||
on_split,
|
||||
raw[AAPL].loc[split_date - trading_day, 'close'] / split_ratio,
|
||||
decimal=2,
|
||||
)
|
||||
|
||||
# length 1 vwap on the day after the split should be the as-traded
|
||||
# close on the split day.
|
||||
after_split = vwaps[1][AAPL].loc[split_date + trading_day]
|
||||
assert_almost_equal(after_split, 93.69999, decimal=2)
|
||||
assert_almost_equal(
|
||||
after_split,
|
||||
raw[AAPL].loc[split_date, 'close'],
|
||||
decimal=2,
|
||||
)
|
||||
|
||||
def initialize(context):
|
||||
context.vwaps = []
|
||||
@@ -195,6 +367,7 @@ class FFCAlgorithmTestCase(TestCase):
|
||||
for asset in assets:
|
||||
computed = factors.loc[asset, key]
|
||||
expected = vwaps[length][asset].loc[today]
|
||||
|
||||
# Only having two places of precision here is a bit
|
||||
# unfortunate.
|
||||
assert_almost_equal(computed, expected, decimal=2)
|
||||
|
||||
@@ -65,7 +65,9 @@ class NumericalExpressionTestCase(TestCase):
|
||||
def check_output(self, expr, expected):
|
||||
result = expr._compute(
|
||||
[self.fake_raw_data[input_] for input_ in expr.inputs],
|
||||
self.mask,
|
||||
self.mask.index,
|
||||
self.mask.columns,
|
||||
self.mask.values,
|
||||
)
|
||||
check_arrays(result, expected)
|
||||
|
||||
|
||||
@@ -19,9 +19,10 @@ from zipline.errors import (
|
||||
TermInputsNotSpecified,
|
||||
WindowLengthNotSpecified,
|
||||
)
|
||||
from zipline.modelling.expression import NUMEXPR_MATH_FUNCS
|
||||
from zipline.modelling.factor import Factor
|
||||
from zipline.modelling.graph import TermGraph
|
||||
from zipline.modelling.expression import NUMEXPR_MATH_FUNCS
|
||||
from zipline.modelling.term import AssetExists, NotSpecified
|
||||
|
||||
|
||||
class SomeDataSet(DataSet):
|
||||
@@ -54,13 +55,16 @@ def gen_equivalent_factors():
|
||||
object.
|
||||
"""
|
||||
yield SomeFactor()
|
||||
yield SomeFactor(inputs=None)
|
||||
yield SomeFactor(inputs=NotSpecified)
|
||||
yield SomeFactor(SomeFactor.inputs)
|
||||
yield SomeFactor(inputs=SomeFactor.inputs)
|
||||
yield SomeFactor([SomeDataSet.foo, SomeDataSet.bar])
|
||||
yield SomeFactor(window_length=SomeFactor.window_length)
|
||||
yield SomeFactor(window_length=None)
|
||||
yield SomeFactor([SomeDataSet.foo, SomeDataSet.bar], window_length=None)
|
||||
yield SomeFactor(window_length=NotSpecified)
|
||||
yield SomeFactor(
|
||||
[SomeDataSet.foo, SomeDataSet.bar],
|
||||
window_length=NotSpecified,
|
||||
)
|
||||
yield SomeFactor(
|
||||
[SomeDataSet.foo, SomeDataSet.bar],
|
||||
window_length=SomeFactor.window_length,
|
||||
@@ -96,9 +100,10 @@ class DependencyResolutionTestCase(TestCase):
|
||||
|
||||
resolution_order = list(graph.ordered())
|
||||
|
||||
self.assertEqual(len(resolution_order), 3)
|
||||
self.assertEqual(len(resolution_order), 4)
|
||||
self.assertIs(resolution_order[0], AssetExists())
|
||||
self.assertEqual(
|
||||
set([resolution_order[0], resolution_order[1]]),
|
||||
set([resolution_order[1], resolution_order[2]]),
|
||||
set([SomeDataSet.foo, SomeDataSet.bar]),
|
||||
)
|
||||
self.assertEqual(resolution_order[-1], SomeFactor())
|
||||
@@ -118,9 +123,14 @@ class DependencyResolutionTestCase(TestCase):
|
||||
|
||||
resolution_order = list(graph.ordered())
|
||||
|
||||
self.assertEqual(len(resolution_order), 3)
|
||||
# SomeFactor, its inputs, and AssetExists()
|
||||
self.assertEqual(len(resolution_order), 4)
|
||||
|
||||
self.assertIs(resolution_order[0], AssetExists())
|
||||
self.assertEqual(graph.extra_rows[AssetExists()], 4)
|
||||
|
||||
self.assertEqual(
|
||||
set([resolution_order[0], resolution_order[1]]),
|
||||
set([resolution_order[1], resolution_order[2]]),
|
||||
set([bar, buzz]),
|
||||
)
|
||||
self.assertEqual(
|
||||
@@ -141,12 +151,14 @@ class DependencyResolutionTestCase(TestCase):
|
||||
resolution_order = list(graph.ordered())
|
||||
|
||||
# bar should only appear once.
|
||||
self.assertEqual(len(resolution_order), 5)
|
||||
self.assertEqual(len(resolution_order), 6)
|
||||
indices = {
|
||||
term: resolution_order.index(term)
|
||||
for term in resolution_order
|
||||
}
|
||||
|
||||
self.assertEqual(indices[AssetExists()], 0)
|
||||
|
||||
# Verify that f1's dependencies will be computed before f1.
|
||||
self.assertLess(indices[SomeDataSet.foo], indices[f1])
|
||||
self.assertLess(indices[SomeDataSet.bar], indices[f1])
|
||||
|
||||
@@ -21,6 +21,8 @@ from nose_parameterized import parameterized
|
||||
from numpy import (
|
||||
arange,
|
||||
datetime64,
|
||||
float64,
|
||||
ones,
|
||||
uint32,
|
||||
)
|
||||
from numpy.testing import (
|
||||
@@ -31,6 +33,7 @@ from pandas import (
|
||||
concat,
|
||||
DataFrame,
|
||||
DatetimeIndex,
|
||||
Int64Index,
|
||||
Timestamp,
|
||||
)
|
||||
from pandas.util.testing import assert_index_equal
|
||||
@@ -206,8 +209,8 @@ class BcolzDailyBarTestCase(TestCase):
|
||||
def _check_read_results(self, columns, assets, start_date, end_date):
|
||||
table = self.writer.write(self.dest, self.trading_days, self.assets)
|
||||
reader = BcolzDailyBarReader(table)
|
||||
results = reader.load_raw_arrays(columns, start_date, end_date, assets)
|
||||
dates = self.trading_days_between(start_date, end_date)
|
||||
results = reader.load_raw_arrays(columns, dates, assets)
|
||||
for column, result in zip(columns, results):
|
||||
assert_array_equal(
|
||||
result,
|
||||
@@ -433,10 +436,14 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
self.assertGreaterEqual(eff_date, asset_start)
|
||||
self.assertLessEqual(eff_date, asset_end)
|
||||
|
||||
def calendar_days_between(self, start_date, end_date):
|
||||
return self.calendar_days[
|
||||
self.calendar_days.slice_indexer(start_date, end_date)
|
||||
]
|
||||
def calendar_days_between(self, start_date, end_date, shift=0):
|
||||
slice_ = self.calendar_days.slice_indexer(start_date, end_date)
|
||||
start = slice_.start + shift
|
||||
stop = slice_.stop + shift
|
||||
if start < 0:
|
||||
raise KeyError(start_date, shift)
|
||||
|
||||
return self.calendar_days[start:stop]
|
||||
|
||||
def expected_adjustments(self, start_date, end_date):
|
||||
price_adjustments = {}
|
||||
@@ -448,23 +455,20 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
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):
|
||||
# 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 adjusments should be applied on the date
|
||||
# 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 - 1,
|
||||
last_row=delta,
|
||||
col=sid - 1,
|
||||
value=ratio,
|
||||
)
|
||||
@@ -474,7 +478,7 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
volume_adjustments.setdefault(delta, []).append(
|
||||
Float64Multiply(
|
||||
first_row=0,
|
||||
last_row=delta - 1,
|
||||
last_row=delta,
|
||||
col=sid - 1,
|
||||
value=1.0 / ratio,
|
||||
)
|
||||
@@ -486,7 +490,7 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
columns = [USEquityPricing.close, USEquityPricing.volume]
|
||||
query_days = self.calendar_days_between(
|
||||
TEST_QUERY_START,
|
||||
TEST_QUERY_STOP
|
||||
TEST_QUERY_STOP,
|
||||
)
|
||||
|
||||
adjustments = reader.load_adjustments(
|
||||
@@ -510,6 +514,13 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
TEST_QUERY_START,
|
||||
TEST_QUERY_STOP
|
||||
)
|
||||
# Our expected results for each day are based on values from the
|
||||
# previous day.
|
||||
shifted_query_days = self.calendar_days_between(
|
||||
TEST_QUERY_START,
|
||||
TEST_QUERY_STOP,
|
||||
shift=-1,
|
||||
)
|
||||
|
||||
adjustments = adjustment_reader.load_adjustments(
|
||||
columns,
|
||||
@@ -526,16 +537,18 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
|
||||
closes, volumes = pricing_loader.load_adjusted_array(
|
||||
columns,
|
||||
DataFrame(True, index=query_days, columns=self.assets),
|
||||
dates=query_days,
|
||||
assets=self.assets,
|
||||
mask=ones((len(query_days), len(self.assets)), dtype=bool),
|
||||
)
|
||||
|
||||
expected_baseline_closes = self.bcolz_writer.expected_values_2d(
|
||||
query_days,
|
||||
shifted_query_days,
|
||||
self.assets,
|
||||
'close',
|
||||
)
|
||||
expected_baseline_volumes = self.bcolz_writer.expected_values_2d(
|
||||
query_days,
|
||||
shifted_query_days,
|
||||
self.assets,
|
||||
'volume',
|
||||
)
|
||||
@@ -562,19 +575,22 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
|
||||
def apply_adjustments(self, dates, assets, baseline_values, adjustments):
|
||||
min_date, max_date = dates[[0, -1]]
|
||||
values = baseline_values.copy()
|
||||
# HACK: Simulate the coercion to float64 we do in adjusted_array. This
|
||||
# should be removed when AdjustedArray properly supports
|
||||
# non-floating-point types.
|
||||
orig_dtype = baseline_values.dtype
|
||||
values = baseline_values.astype(float64).copy()
|
||||
for eff_date_secs, ratio, sid in adjustments.itertuples(index=False):
|
||||
eff_date = seconds_to_timestamp(eff_date_secs)
|
||||
if eff_date < min_date or eff_date > max_date:
|
||||
# Don't apply adjustments that aren't in the current date range.
|
||||
if eff_date not in dates:
|
||||
continue
|
||||
eff_date_loc = dates.get_loc(eff_date)
|
||||
asset_col = assets.get_loc(sid)
|
||||
# Apply ratio multiplicatively to the asset column on all rows
|
||||
# **strictly less** than the adjustment effective date. Note that
|
||||
# this will be a no-op in the case that the effective date is the
|
||||
# first entry in dates.
|
||||
values[:eff_date_loc, asset_col] *= ratio
|
||||
return values
|
||||
# Apply ratio multiplicatively to the asset column on all rows less
|
||||
# than or equal adjustment effective date.
|
||||
values[:eff_date_loc + 1, asset_col] *= ratio
|
||||
return values.astype(orig_dtype)
|
||||
|
||||
def test_read_with_adjustments(self):
|
||||
columns = [USEquityPricing.high, USEquityPricing.volume]
|
||||
@@ -582,6 +598,13 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
TEST_QUERY_START,
|
||||
TEST_QUERY_STOP
|
||||
)
|
||||
# Our expected results for each day are based on values from the
|
||||
# previous day.
|
||||
shifted_query_days = self.calendar_days_between(
|
||||
TEST_QUERY_START,
|
||||
TEST_QUERY_STOP,
|
||||
shift=-1,
|
||||
)
|
||||
|
||||
baseline_reader = BcolzDailyBarReader(self.bcolz_path)
|
||||
adjustment_reader = SQLiteAdjustmentReader(self.db_path)
|
||||
@@ -590,18 +613,20 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
adjustment_reader,
|
||||
)
|
||||
|
||||
closes, volumes = pricing_loader.load_adjusted_array(
|
||||
highs, volumes = pricing_loader.load_adjusted_array(
|
||||
columns,
|
||||
DataFrame(True, index=query_days, columns=arange(1, 7)),
|
||||
dates=query_days,
|
||||
assets=Int64Index(arange(1, 7)),
|
||||
mask=ones((len(query_days), 6), dtype=bool),
|
||||
)
|
||||
|
||||
expected_baseline_highs = self.bcolz_writer.expected_values_2d(
|
||||
query_days,
|
||||
shifted_query_days,
|
||||
self.assets,
|
||||
'high',
|
||||
)
|
||||
expected_baseline_volumes = self.bcolz_writer.expected_values_2d(
|
||||
query_days,
|
||||
shifted_query_days,
|
||||
self.assets,
|
||||
'volume',
|
||||
)
|
||||
@@ -609,7 +634,7 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
# At each point in time, the AdjustedArrays should yield the baseline
|
||||
# with all adjustments up to that date applied.
|
||||
for windowlen in range(1, len(query_days) + 1):
|
||||
for offset, window in enumerate(closes.traverse(windowlen)):
|
||||
for offset, window in enumerate(highs.traverse(windowlen)):
|
||||
baseline = expected_baseline_highs[offset:offset + windowlen]
|
||||
baseline_dates = query_days[offset:offset + windowlen]
|
||||
expected_adjusted_highs = self.apply_adjustments(
|
||||
@@ -627,6 +652,7 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
# Apply only splits and invert the ratio.
|
||||
adjustments = SPLITS.copy()
|
||||
adjustments.ratio = 1 / adjustments.ratio
|
||||
|
||||
expected_adjusted_volumes = self.apply_adjustments(
|
||||
baseline_dates,
|
||||
self.assets,
|
||||
@@ -641,6 +667,6 @@ class USEquityPricingLoaderTestCase(TestCase):
|
||||
|
||||
# Verify that we checked up to the longest possible window.
|
||||
with self.assertRaises(WindowLengthTooLong):
|
||||
closes.traverse(windowlen + 1)
|
||||
highs.traverse(windowlen + 1)
|
||||
with self.assertRaises(WindowLengthTooLong):
|
||||
volumes.traverse(windowlen + 1)
|
||||
|
||||
+22
-7
@@ -653,7 +653,12 @@ class AssetFinderTestCase(TestCase):
|
||||
)
|
||||
|
||||
for dates in all_subindices(all_dates):
|
||||
expected_mask = full(
|
||||
expected_with_start_raw = full(
|
||||
shape=(len(dates), num_assets),
|
||||
fill_value=False,
|
||||
dtype=bool,
|
||||
)
|
||||
expected_no_start_raw = full(
|
||||
shape=(len(dates), num_assets),
|
||||
fill_value=False,
|
||||
dtype=bool,
|
||||
@@ -662,18 +667,28 @@ class AssetFinderTestCase(TestCase):
|
||||
for i, date in enumerate(dates):
|
||||
it = frame[['start_date', 'end_date']].itertuples()
|
||||
for j, start, end in it:
|
||||
# This way of doing the checks is redundant, but very
|
||||
# clear.
|
||||
if start <= date <= end:
|
||||
expected_mask[i, j] = True
|
||||
expected_with_start_raw[i, j] = True
|
||||
if start < date:
|
||||
expected_no_start_raw[i, j] = True
|
||||
|
||||
# Filter out columns with all-empty columns.
|
||||
expected_result = pd.DataFrame(
|
||||
data=expected_mask,
|
||||
expected_with_start = pd.DataFrame(
|
||||
data=expected_with_start_raw,
|
||||
index=dates,
|
||||
columns=frame.index.values,
|
||||
)
|
||||
result = finder.lifetimes(dates, include_start_date=True)
|
||||
assert_frame_equal(result, expected_with_start)
|
||||
|
||||
actual_result = finder.lifetimes(dates)
|
||||
assert_frame_equal(actual_result, expected_result)
|
||||
expected_no_start = pd.DataFrame(
|
||||
data=expected_no_start_raw,
|
||||
index=dates,
|
||||
columns=frame.index.values,
|
||||
)
|
||||
result = finder.lifetimes(dates, include_start_date=False)
|
||||
assert_frame_equal(result, expected_no_start)
|
||||
|
||||
def test_sids(self):
|
||||
# Ensure that the sids property of the AssetFinder is functioning
|
||||
|
||||
+12
-1
@@ -1363,13 +1363,24 @@ class TradingAlgorithm(object):
|
||||
|
||||
def compute_factor_matrix(self, start_date):
|
||||
"""
|
||||
Compute a factor matrix starting at start_date.
|
||||
Compute a factor matrix containing at least the data necessary to
|
||||
provide values for `start_date`.
|
||||
|
||||
Loads a factor matrix with data extending from `start_date` until a
|
||||
year from `start_date`, or until the end of the simulation.
|
||||
"""
|
||||
days = self.trading_environment.trading_days
|
||||
|
||||
# Load data starting from the previous trading day...
|
||||
start_date_loc = days.get_loc(start_date)
|
||||
|
||||
# ...continuing until either the day before the simulation end, or
|
||||
# until 252 days of data have been loaded. 252 is a totally arbitrary
|
||||
# choice that seemed reasonable based on napkin math.
|
||||
sim_end = self.sim_params.last_close.normalize()
|
||||
end_loc = min(start_date_loc + 252, days.get_loc(sim_end))
|
||||
end_date = days[end_loc]
|
||||
|
||||
return self.engine.factor_matrix(
|
||||
self._all_terms(),
|
||||
start_date,
|
||||
|
||||
@@ -625,7 +625,7 @@ class AssetFinder(object):
|
||||
('end', '<i8'),
|
||||
])
|
||||
|
||||
def lifetimes(self, dates):
|
||||
def lifetimes(self, dates, include_start_date):
|
||||
"""
|
||||
Compute a DataFrame representing asset lifetimes for the specified date
|
||||
range.
|
||||
@@ -634,17 +634,28 @@ class AssetFinder(object):
|
||||
----------
|
||||
dates : pd.DatetimeIndex
|
||||
The dates for which to compute lifetimes.
|
||||
include_start_date : bool
|
||||
Whether or not to count the asset as alive on its start_date.
|
||||
|
||||
This is useful in a backtesting context where `lifetimes` is being
|
||||
used to signify "do I have data for this asset as of the morning of
|
||||
this date?" For many financial metrics, (e.g. daily close), data
|
||||
isn't available for an asset until the end of the asset's first
|
||||
day.
|
||||
|
||||
Returns
|
||||
-------
|
||||
lifetimes : pd.DataFrame
|
||||
A frame of dtype bool with `dates` as index and an Int64Index of
|
||||
assets as columns. The value at `lifetimes.loc[date, asset]` will
|
||||
be True iff `asset` existed on `data`.
|
||||
be True iff `asset` existed on `date`. If `include_start_date` is
|
||||
False, then lifetimes.loc[date, asset] will be false when date ==
|
||||
asset.start_date.
|
||||
|
||||
See Also
|
||||
--------
|
||||
numpy.putmask
|
||||
zipline.modelling.engine.SimpleFFCEngine._compute_root_mask
|
||||
"""
|
||||
# This is a less than ideal place to do this, because if someone adds
|
||||
# assets to the finder after we've touched lifetimes we won't have
|
||||
@@ -655,7 +666,12 @@ class AssetFinder(object):
|
||||
lifetimes = self._asset_lifetimes
|
||||
|
||||
raw_dates = dates.asi8[:, None]
|
||||
mask = (lifetimes.start <= raw_dates) & (raw_dates <= lifetimes.end)
|
||||
if include_start_date:
|
||||
mask = lifetimes.start <= raw_dates
|
||||
else:
|
||||
mask = lifetimes.start < raw_dates
|
||||
mask &= (raw_dates <= lifetimes.end)
|
||||
|
||||
return pd.DataFrame(mask, index=dates, columns=lifetimes.sid)
|
||||
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from six import (
|
||||
)
|
||||
|
||||
from zipline.modelling.term import Term
|
||||
from zipline.modelling.factor import Latest
|
||||
|
||||
|
||||
class Column(object):
|
||||
@@ -68,6 +69,10 @@ class BoundColumn(Term):
|
||||
"""
|
||||
return '.'.join([self.dataset.__name__, self.name])
|
||||
|
||||
@property
|
||||
def latest(self):
|
||||
return Latest(inputs=(self,), dtype=self.dtype)
|
||||
|
||||
def __repr__(self):
|
||||
return "{qualname}::{dtype}".format(
|
||||
qualname=self.qualname,
|
||||
|
||||
@@ -47,10 +47,12 @@ class DataFrameFFCLoader(FFCLoader):
|
||||
----------
|
||||
column : zipline.data.dataset.BoundColumn
|
||||
The column whose data is loadable by this loader.
|
||||
|
||||
baseline : pandas.DataFrame
|
||||
A DataFrame with index of type DatetimeIndex and columns of type
|
||||
Int64Index.
|
||||
Int64Index. Dates should be labelled with the first date on which a
|
||||
value would be **available** to an algorithm. This means that OHLCV
|
||||
data should generally be shifted back by a trading day before being
|
||||
supplied to this class.
|
||||
|
||||
adjustments : pandas.DataFrame, default=None
|
||||
A DataFrame with the following columns:
|
||||
@@ -156,7 +158,7 @@ class DataFrameFFCLoader(FFCLoader):
|
||||
)
|
||||
return out
|
||||
|
||||
def load_adjusted_array(self, columns, mask):
|
||||
def load_adjusted_array(self, columns, dates, assets, mask):
|
||||
"""
|
||||
Load data from our stored baseline.
|
||||
"""
|
||||
@@ -167,8 +169,6 @@ class DataFrameFFCLoader(FFCLoader):
|
||||
elif columns[0] != self.column:
|
||||
raise ValueError("Can't load unknown column %s" % columns[0])
|
||||
|
||||
dates, assets, mask_values = mask.index, mask.columns, mask.values
|
||||
|
||||
date_indexer = self.dates.get_indexer(dates)
|
||||
assets_indexer = self.assets.get_indexer(assets)
|
||||
|
||||
@@ -180,6 +180,6 @@ class DataFrameFFCLoader(FFCLoader):
|
||||
# Pull out requested columns/rows from our baseline data.
|
||||
data=self.baseline[ix_(date_indexer, assets_indexer)],
|
||||
# Mask out requested columns/rows that didnt match.
|
||||
mask=(good_assets & good_dates[:, None]) & mask_values,
|
||||
mask=(good_assets & good_dates[:, None]) & mask,
|
||||
adjustments=self.format_adjustments(dates, assets),
|
||||
)]
|
||||
|
||||
@@ -0,0 +1,282 @@
|
||||
#
|
||||
# Copyright 2015 Quantopian, Inc.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from cpython cimport (
|
||||
PyDict_Contains,
|
||||
PySet_Add,
|
||||
)
|
||||
|
||||
from numpy import (
|
||||
uint32,
|
||||
zeros,
|
||||
)
|
||||
from pandas import Timestamp
|
||||
|
||||
ctypedef object Timestamp_t
|
||||
ctypedef object DatetimeIndex_t
|
||||
ctypedef object Int64Index_t
|
||||
|
||||
from zipline.lib.adjustment import Float64Multiply
|
||||
|
||||
_SID_QUERY_TEMPLATE = """
|
||||
SELECT DISTINCT sid FROM {0}
|
||||
WHERE effective_date >= ? AND effective_date <= ?
|
||||
"""
|
||||
cdef dict SID_QUERIES = {
|
||||
tablename: _SID_QUERY_TEMPLATE.format(tablename)
|
||||
for tablename in ('splits', 'dividends', 'mergers')
|
||||
}
|
||||
|
||||
ADJ_QUERY_TEMPLATE = """
|
||||
SELECT sid, ratio, effective_date
|
||||
FROM {0}
|
||||
WHERE sid IN ({1}) AND effective_date >= {2} AND effective_date <= {3}
|
||||
"""
|
||||
|
||||
cdef int SQLITE_MAX_IN_STATEMENT = 999
|
||||
EPOCH = Timestamp(0, tz='UTC')
|
||||
|
||||
cdef set _get_sids_from_table(object db,
|
||||
str tablename,
|
||||
int start_date,
|
||||
int end_date):
|
||||
"""
|
||||
Get the unique sids for all adjustments between start_date and end_date
|
||||
from table `tablename`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
db : sqlite3.connection
|
||||
tablename : str
|
||||
start_date : int (seconds since epoch)
|
||||
end_date : int (seconds since epoch)
|
||||
|
||||
Returns
|
||||
-------
|
||||
sids : set
|
||||
Set of sets
|
||||
"""
|
||||
|
||||
cdef object cursor = db.execute(
|
||||
SID_QUERIES[tablename],
|
||||
(start_date, end_date),
|
||||
)
|
||||
cdef set out = set()
|
||||
cdef tuple result
|
||||
for result in cursor.fetchall():
|
||||
PySet_Add(out, result[0])
|
||||
return out
|
||||
|
||||
|
||||
cdef set _get_split_sids(object db, int start_date, int end_date):
|
||||
return _get_sids_from_table(db, 'splits', start_date, end_date)
|
||||
|
||||
|
||||
cdef set _get_merger_sids(object db, int start_date, int end_date):
|
||||
return _get_sids_from_table(db, 'mergers', start_date, end_date)
|
||||
|
||||
|
||||
cdef set _get_dividend_sids(object db, int start_date, int end_date):
|
||||
return _get_sids_from_table(db, 'dividends', start_date, end_date)
|
||||
|
||||
|
||||
cdef _adjustments(object adjustments_db,
|
||||
set split_sids,
|
||||
set merger_sids,
|
||||
set dividends_sids,
|
||||
int start_date,
|
||||
int end_date,
|
||||
Int64Index_t assets):
|
||||
|
||||
c = adjustments_db.cursor()
|
||||
|
||||
splits_to_query = [str(a) for a in assets if a in split_sids]
|
||||
splits_results = []
|
||||
while splits_to_query:
|
||||
query_len = min(len(splits_to_query), SQLITE_MAX_IN_STATEMENT)
|
||||
query_assets = splits_to_query[:query_len]
|
||||
t= [str(a) for a in query_assets]
|
||||
statement = ADJ_QUERY_TEMPLATE.format('splits',
|
||||
",".join(['?' for _ in query_assets]), start_date, end_date)
|
||||
c.execute(statement, t)
|
||||
splits_to_query = splits_to_query[query_len:]
|
||||
splits_results.extend(c.fetchall())
|
||||
|
||||
mergers_to_query = [str(a) for a in assets if a in merger_sids]
|
||||
mergers_results = []
|
||||
while mergers_to_query:
|
||||
query_len = min(len(mergers_to_query), SQLITE_MAX_IN_STATEMENT)
|
||||
query_assets = mergers_to_query[:query_len]
|
||||
t= [str(a) for a in query_assets]
|
||||
statement = ADJ_QUERY_TEMPLATE.format('mergers',
|
||||
",".join(['?' for _ in query_assets]), start_date, end_date)
|
||||
c.execute(statement, t)
|
||||
mergers_to_query = mergers_to_query[query_len:]
|
||||
mergers_results.extend(c.fetchall())
|
||||
|
||||
dividends_to_query = [str(a) for a in assets if a in dividends_sids]
|
||||
dividends_results = []
|
||||
while dividends_to_query:
|
||||
query_len = min(len(dividends_to_query), SQLITE_MAX_IN_STATEMENT)
|
||||
query_assets = dividends_to_query[:query_len]
|
||||
t= [str(a) for a in query_assets]
|
||||
statement = ADJ_QUERY_TEMPLATE.format('dividends',
|
||||
",".join(['?' for _ in query_assets]), start_date, end_date)
|
||||
c.execute(statement, t)
|
||||
dividends_to_query = dividends_to_query[query_len:]
|
||||
dividends_results.extend(c.fetchall())
|
||||
|
||||
return splits_results, mergers_results, dividends_results
|
||||
|
||||
|
||||
cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
list columns,
|
||||
DatetimeIndex_t dates,
|
||||
Int64Index_t assets):
|
||||
"""
|
||||
Load a dictionary of Adjustment objects from adjustments_db
|
||||
|
||||
Parameters
|
||||
----------
|
||||
adjustments_db : sqlite3.Connection
|
||||
Connection to a sqlite3 table in the format written by
|
||||
SQLiteAdjustmentWriter.
|
||||
columns : list[str]
|
||||
List of column names for which adjustments are needed.
|
||||
dates : pd.DatetimeIndex
|
||||
Dates for which adjustments are needed
|
||||
assets : pd.Int64Index
|
||||
Assets for which adjustments are needed.
|
||||
"""
|
||||
|
||||
cdef int start_date = int((dates[0] - EPOCH).total_seconds())
|
||||
cdef int end_date = int((dates[-1] - EPOCH).total_seconds())
|
||||
|
||||
cdef set split_sids = _get_split_sids(
|
||||
adjustments_db,
|
||||
start_date,
|
||||
end_date,
|
||||
)
|
||||
cdef set merger_sids = _get_merger_sids(
|
||||
adjustments_db,
|
||||
start_date,
|
||||
end_date,
|
||||
)
|
||||
cdef set dividend_sids = _get_dividend_sids(
|
||||
adjustments_db,
|
||||
start_date,
|
||||
end_date,
|
||||
)
|
||||
|
||||
cdef:
|
||||
list splits, mergers, dividends
|
||||
splits, mergers, dividends = _adjustments(
|
||||
adjustments_db,
|
||||
split_sids,
|
||||
merger_sids,
|
||||
dividend_sids,
|
||||
start_date,
|
||||
end_date,
|
||||
assets,
|
||||
)
|
||||
|
||||
cdef list results = [{} for column in columns]
|
||||
cdef dict asset_ixs = {} # Cache sid lookups here.
|
||||
cdef:
|
||||
int sid
|
||||
double ratio
|
||||
int eff_date
|
||||
int date_loc
|
||||
Py_ssize_t asset_ix
|
||||
int i
|
||||
dict col_adjustments
|
||||
|
||||
# splits affect prices and volumes, volumes is the inverse
|
||||
for sid, ratio, eff_date in splits:
|
||||
if eff_date < start_date:
|
||||
continue
|
||||
date_loc = dates.get_loc(
|
||||
Timestamp(eff_date, unit='s', tz='UTC'),
|
||||
# Get the first date **on or after** the effective date.
|
||||
method='bfill',
|
||||
)
|
||||
|
||||
if not PyDict_Contains(asset_ixs, sid):
|
||||
asset_ixs[sid] = assets.get_loc(sid)
|
||||
asset_ix = asset_ixs[sid]
|
||||
|
||||
price_adj = Float64Multiply(0, date_loc, asset_ix, ratio)
|
||||
for i, column in enumerate(columns):
|
||||
col_adjustments = results[i]
|
||||
if column != 'volume':
|
||||
try:
|
||||
col_adjustments[date_loc].append(price_adj)
|
||||
except KeyError:
|
||||
col_adjustments[date_loc] = [price_adj]
|
||||
else:
|
||||
volume_adj = Float64Multiply(
|
||||
0, date_loc, asset_ix, 1.0 / ratio
|
||||
)
|
||||
try:
|
||||
col_adjustments[date_loc].append(volume_adj)
|
||||
except KeyError:
|
||||
col_adjustments[date_loc] = [volume_adj]
|
||||
|
||||
# mergers affect prices only
|
||||
for sid, ratio, eff_date in mergers:
|
||||
if eff_date < start_date:
|
||||
continue
|
||||
date_loc = dates.get_loc(
|
||||
Timestamp(eff_date, unit='s', tz='UTC'),
|
||||
# Get the first date **on or after** the effective date.
|
||||
method='bfill',
|
||||
)
|
||||
|
||||
if not PyDict_Contains(asset_ixs, sid):
|
||||
asset_ixs[sid] = assets.get_loc(sid)
|
||||
asset_ix = asset_ixs[sid]
|
||||
|
||||
adj = Float64Multiply(0, date_loc, asset_ix, ratio)
|
||||
for i, column in enumerate(columns):
|
||||
col_adjustments = results[i]
|
||||
if column != 'volume':
|
||||
try:
|
||||
col_adjustments[date_loc].append(adj)
|
||||
except KeyError:
|
||||
col_adjustments[date_loc] = [adj]
|
||||
|
||||
# dividends affect prices only
|
||||
for sid, ratio, eff_date in dividends:
|
||||
if eff_date < start_date:
|
||||
continue
|
||||
date_loc = dates.get_loc(
|
||||
Timestamp(eff_date, unit='s', tz='UTC'),
|
||||
# Get the first date **on or after** the effective date.
|
||||
method='bfill',
|
||||
)
|
||||
|
||||
if not PyDict_Contains(asset_ixs, sid):
|
||||
asset_ixs[sid] = assets.get_loc(sid)
|
||||
asset_ix = asset_ixs[sid]
|
||||
|
||||
adj = Float64Multiply(0, date_loc, asset_ix, ratio)
|
||||
for i, column in enumerate(columns):
|
||||
col_adjustments = results[i]
|
||||
if column != 'volume':
|
||||
try:
|
||||
col_adjustments[date_loc].append(adj)
|
||||
except KeyError:
|
||||
col_adjustments[date_loc] = [adj]
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,196 @@
|
||||
#
|
||||
# Copyright 2015 Quantopian, Inc.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import bcolz
|
||||
cimport cython
|
||||
|
||||
from numpy import (
|
||||
array,
|
||||
float64,
|
||||
intp,
|
||||
uint32,
|
||||
zeros,
|
||||
)
|
||||
from numpy cimport (
|
||||
float64_t,
|
||||
intp_t,
|
||||
ndarray,
|
||||
uint32_t,
|
||||
uint8_t,
|
||||
)
|
||||
from numpy.math cimport NAN
|
||||
|
||||
ctypedef object ctable_t
|
||||
ctypedef object Timestamp_t
|
||||
ctypedef object DatetimeIndex_t
|
||||
ctypedef object Int64Index_t
|
||||
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cpdef _compute_row_slices(dict asset_starts_absolute,
|
||||
dict asset_ends_absolute,
|
||||
dict asset_starts_calendar,
|
||||
intp_t query_start,
|
||||
intp_t query_end,
|
||||
Int64Index_t requested_assets):
|
||||
"""
|
||||
Core indexing functionality for loading raw data from bcolz.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
asset_starts_absolute : dict
|
||||
Dictionary containing the index of the first row of each asset in the
|
||||
bcolz file from which we will query.
|
||||
|
||||
asset_ends_absolute : dict
|
||||
Dictionary containing the index of the last row of each asset in the
|
||||
bcolz file from which we will query.
|
||||
|
||||
asset_starts_calendar : dict
|
||||
Dictionary containing the index of in our calendar corresponding to the
|
||||
start date of each asset
|
||||
|
||||
query_start : intp
|
||||
query_end : intp
|
||||
Start and end indices in our calendar of the dates for which we're
|
||||
querying.
|
||||
|
||||
requested_assets : pandas.Int64Index
|
||||
The assets for which we want to load data.
|
||||
|
||||
For each asset in requested assets, computes three values:
|
||||
1.) The index in the raw bcolz data of first row to load.
|
||||
2.) The index in the raw bcolz data of the last row to load.
|
||||
3.) The index in the dates of our query corresponding to the first row for
|
||||
each asset. This is non-zero iff the asset's lifetime begins partway
|
||||
through the requested query dates.
|
||||
|
||||
Returns
|
||||
-------
|
||||
first_rows, last_rows, offsets : 3-tuple of ndarrays
|
||||
"""
|
||||
cdef:
|
||||
intp_t nassets = len(requested_assets)
|
||||
|
||||
# For each sid, we need to compute the following:
|
||||
ndarray[dtype=intp_t, ndim=1] first_row_a = zeros(nassets, dtype=intp)
|
||||
ndarray[dtype=intp_t, ndim=1] last_row_a = zeros(nassets, dtype=intp)
|
||||
ndarray[dtype=intp_t, ndim=1] offset_a = zeros(nassets, dtype=intp)
|
||||
|
||||
# Loop variables.
|
||||
intp_t i
|
||||
intp_t asset
|
||||
intp_t asset_start_data
|
||||
intp_t asset_end_data
|
||||
intp_t asset_start_calendar
|
||||
intp_t asset_end_calendar
|
||||
|
||||
for i, asset in enumerate(requested_assets):
|
||||
asset_start_data = asset_starts_absolute[asset]
|
||||
asset_end_data = asset_ends_absolute[asset]
|
||||
asset_start_calendar = asset_starts_calendar[asset]
|
||||
asset_end_calendar = (
|
||||
asset_start_calendar + (asset_end_data - asset_start_data)
|
||||
)
|
||||
|
||||
# If the asset started during the query, then start with the asset's
|
||||
# first row.
|
||||
# Otherwise start with the asset's first row + the number of rows
|
||||
# before the query on which the asset existed.
|
||||
first_row_a[i] = (
|
||||
asset_start_data + max(0, (query_start - asset_start_calendar))
|
||||
)
|
||||
# If the asset ended during the query, the end with the asset's last
|
||||
# row.
|
||||
# Otherwise, end with the asset's last row minus the number of rows
|
||||
# after the query for which the asset
|
||||
last_row_a[i] = (
|
||||
asset_end_data - max(0, asset_end_calendar - query_end)
|
||||
)
|
||||
# If the asset existed on or before the query, no offset.
|
||||
# Otherwise, offset by the number of rows in the query in which the
|
||||
# asset did not yet exist.
|
||||
offset_a[i] = max(0, asset_start_calendar - query_start)
|
||||
|
||||
return first_row_a, last_row_a, offset_a
|
||||
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cpdef _read_bcolz_data(ctable_t table,
|
||||
tuple shape,
|
||||
list columns,
|
||||
intp_t[:] first_rows,
|
||||
intp_t[:] last_rows,
|
||||
intp_t[:] offsets):
|
||||
"""
|
||||
Load raw bcolz data for the given columns and indices.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
table : bcolz.ctable
|
||||
The table from which to read.
|
||||
shape : tuple (length 2)
|
||||
The shape of the expected output arrays.
|
||||
columns : list[str]
|
||||
List of column names to read.
|
||||
|
||||
first_rows : ndarray[intp]
|
||||
last_rows : ndarray[intp]
|
||||
offsets : ndarray[intp
|
||||
Arrays in the format returned by _compute_row_slices.
|
||||
|
||||
Returns
|
||||
-------
|
||||
results : list of ndarray
|
||||
A 2D array of shape `shape` for each column in `columns`.
|
||||
"""
|
||||
cdef:
|
||||
int nassets
|
||||
str column_name
|
||||
ndarray[dtype=uint32_t, ndim=1] raw_data
|
||||
ndarray[dtype=uint32_t, ndim=2] outbuf
|
||||
ndarray[dtype=uint8_t, ndim=2, cast=True] where_nan
|
||||
ndarray[dtype=float64_t, ndim=2] outbuf_as_float
|
||||
intp_t asset
|
||||
intp_t out_idx
|
||||
intp_t raw_idx
|
||||
intp_t first_row
|
||||
intp_t last_row
|
||||
intp_t offset
|
||||
list results = []
|
||||
|
||||
nassets = shape[1]
|
||||
if not nassets== len(first_rows) == len(last_rows) == len(offsets):
|
||||
raise ValueError("Incompatible index arrays.")
|
||||
|
||||
for column_name in columns:
|
||||
raw_data = table[column_name][:]
|
||||
outbuf = zeros(shape=shape, dtype=uint32)
|
||||
for asset in range(nassets):
|
||||
first_row = first_rows[asset]
|
||||
last_row = last_rows[asset]
|
||||
offset = offsets[asset]
|
||||
for out_idx, raw_idx in enumerate(range(first_row, last_row + 1)):
|
||||
outbuf[out_idx + offset, asset] = raw_data[raw_idx]
|
||||
|
||||
if column_name in {'open', 'high', 'low', 'close'}:
|
||||
where_nan = (outbuf == 0)
|
||||
outbuf_as_float = outbuf.astype(float64) * .001
|
||||
outbuf_as_float[where_nan] = NAN
|
||||
results.append(outbuf_as_float)
|
||||
else:
|
||||
results.append(outbuf)
|
||||
return results
|
||||
@@ -17,26 +17,12 @@ from cpython cimport (
|
||||
PySet_Add,
|
||||
)
|
||||
|
||||
import bcolz
|
||||
cimport cython
|
||||
from numpy import (
|
||||
array,
|
||||
float64,
|
||||
intp,
|
||||
uint32,
|
||||
zeros,
|
||||
)
|
||||
from numpy cimport (
|
||||
float64_t,
|
||||
intp_t,
|
||||
ndarray,
|
||||
uint32_t,
|
||||
uint8_t,
|
||||
)
|
||||
from numpy.math cimport NAN
|
||||
from pandas import Timestamp
|
||||
|
||||
ctypedef object ctable_t
|
||||
ctypedef object Timestamp_t
|
||||
ctypedef object DatetimeIndex_t
|
||||
ctypedef object Int64Index_t
|
||||
@@ -212,7 +198,6 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
double ratio
|
||||
int eff_date
|
||||
int date_loc
|
||||
int last_row
|
||||
Py_ssize_t asset_ix
|
||||
int i
|
||||
dict col_adjustments
|
||||
@@ -224,15 +209,12 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
# Get the first date **on or after** the effective date.
|
||||
method='bfill',
|
||||
)
|
||||
last_row = date_loc - 1
|
||||
if last_row < 0:
|
||||
continue
|
||||
|
||||
if not PyDict_Contains(asset_ixs, sid):
|
||||
asset_ixs[sid] = assets.get_loc(sid)
|
||||
asset_ix = asset_ixs[sid]
|
||||
|
||||
price_adj = Float64Multiply(0, last_row, asset_ix, ratio)
|
||||
price_adj = Float64Multiply(0, date_loc, asset_ix, ratio)
|
||||
for i, column in enumerate(columns):
|
||||
col_adjustments = results[i]
|
||||
if column != 'volume':
|
||||
@@ -242,7 +224,7 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
col_adjustments[date_loc] = [price_adj]
|
||||
else:
|
||||
volume_adj = Float64Multiply(
|
||||
0, last_row, asset_ix, 1.0 / ratio
|
||||
0, date_loc, asset_ix, 1.0 / ratio
|
||||
)
|
||||
try:
|
||||
col_adjustments[date_loc].append(volume_adj)
|
||||
@@ -256,15 +238,12 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
# Get the first date **on or after** the effective date.
|
||||
method='bfill',
|
||||
)
|
||||
last_row = date_loc - 1
|
||||
if last_row < 0:
|
||||
continue
|
||||
|
||||
if not PyDict_Contains(asset_ixs, sid):
|
||||
asset_ixs[sid] = assets.get_loc(sid)
|
||||
asset_ix = asset_ixs[sid]
|
||||
|
||||
adj = Float64Multiply(0, last_row, asset_ix, ratio)
|
||||
adj = Float64Multiply(0, date_loc, asset_ix, ratio)
|
||||
for i, column in enumerate(columns):
|
||||
col_adjustments = results[i]
|
||||
if column != 'volume':
|
||||
@@ -280,9 +259,6 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
# Get the first date **on or after** the effective date.
|
||||
method='bfill',
|
||||
)
|
||||
last_row = date_loc - 1
|
||||
if last_row <= 0:
|
||||
continue
|
||||
|
||||
if not PyDict_Contains(asset_ixs, sid):
|
||||
asset_ixs[sid] = assets.get_loc(sid)
|
||||
@@ -298,162 +274,3 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection
|
||||
col_adjustments[date_loc] = [adj]
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cpdef _compute_row_slices(dict asset_starts_absolute,
|
||||
dict asset_ends_absolute,
|
||||
dict asset_starts_calendar,
|
||||
intp_t query_start,
|
||||
intp_t query_end,
|
||||
Int64Index_t requested_assets):
|
||||
"""
|
||||
Core indexing functionality for loading raw data from bcolz.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
asset_starts_absolute : dict
|
||||
Dictionary containing the index of the first row of each asset in the
|
||||
bcolz file from which we will query.
|
||||
|
||||
asset_ends_absolute : dict
|
||||
Dictionary containing the index of the last row of each asset in the
|
||||
bcolz file from which we will query.
|
||||
|
||||
asset_starts_calendar : dict
|
||||
Dictionary containing the index of in our calendar corresponding to the
|
||||
start date of each asset
|
||||
|
||||
query_start : intp
|
||||
query_end : intp
|
||||
Start and end indices in our calendar of the dates for which we're
|
||||
querying.
|
||||
|
||||
requested_assets : pandas.Int64Index
|
||||
The assets for which we want to load data.
|
||||
|
||||
For each asset in requested assets, computes three values:
|
||||
1.) The index in the raw bcolz data of first row to load.
|
||||
2.) The index in the raw bcolz data of the last row to load.
|
||||
3.) The index in the dates of our query corresponding to the first row for
|
||||
each asset. This is non-zero iff the asset's lifetime begins partway
|
||||
through the requested query dates.
|
||||
|
||||
Returns
|
||||
-------
|
||||
first_rows, last_rows, offsets : 3-tuple of ndarrays
|
||||
"""
|
||||
cdef:
|
||||
intp_t nassets = len(requested_assets)
|
||||
|
||||
# For each sid, we need to compute the following:
|
||||
ndarray[dtype=intp_t, ndim=1] first_row_a = zeros(nassets, dtype=intp)
|
||||
ndarray[dtype=intp_t, ndim=1] last_row_a = zeros(nassets, dtype=intp)
|
||||
ndarray[dtype=intp_t, ndim=1] offset_a = zeros(nassets, dtype=intp)
|
||||
|
||||
# Loop variables.
|
||||
intp_t i
|
||||
intp_t asset
|
||||
intp_t asset_start_data
|
||||
intp_t asset_end_data
|
||||
intp_t asset_start_calendar
|
||||
intp_t asset_end_calendar
|
||||
|
||||
for i, asset in enumerate(requested_assets):
|
||||
asset_start_data = asset_starts_absolute[asset]
|
||||
asset_end_data = asset_ends_absolute[asset]
|
||||
asset_start_calendar = asset_starts_calendar[asset]
|
||||
asset_end_calendar = (
|
||||
asset_start_calendar + (asset_end_data - asset_start_data)
|
||||
)
|
||||
|
||||
# If the asset started during the query, then start with the asset's
|
||||
# first row.
|
||||
# Otherwise start with the asset's first row + the number of rows
|
||||
# before the query on which the asset existed.
|
||||
first_row_a[i] = (
|
||||
asset_start_data + max(0, (query_start - asset_start_calendar))
|
||||
)
|
||||
# If the asset ended during the query, the end with the asset's last
|
||||
# row.
|
||||
# Otherwise, end with the asset's last row minus the number of rows
|
||||
# after the query for which the asset
|
||||
last_row_a[i] = (
|
||||
asset_end_data - max(0, asset_end_calendar - query_end)
|
||||
)
|
||||
# If the asset existed on or before the query, no offset.
|
||||
# Otherwise, offset by the number of rows in the query in which the
|
||||
# asset did not yet exist.
|
||||
offset_a[i] = max(0, asset_start_calendar - query_start)
|
||||
|
||||
return first_row_a, last_row_a, offset_a
|
||||
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cpdef _read_bcolz_data(ctable_t table,
|
||||
tuple shape,
|
||||
list columns,
|
||||
intp_t[:] first_rows,
|
||||
intp_t[:] last_rows,
|
||||
intp_t[:] offsets):
|
||||
"""
|
||||
Load raw bcolz data for the given columns and indices.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
table : bcolz.ctable
|
||||
The table from which to read.
|
||||
shape : tuple (length 2)
|
||||
The shape of the expected output arrays.
|
||||
columns : list[str]
|
||||
List of column names to read.
|
||||
|
||||
first_rows : ndarray[intp]
|
||||
last_rows : ndarray[intp]
|
||||
offsets : ndarray[intp
|
||||
Arrays in the format returned by _compute_row_slices.
|
||||
|
||||
Returns
|
||||
-------
|
||||
results : list of ndarray
|
||||
A 2D array of shape `shape` for each column in `columns`.
|
||||
"""
|
||||
cdef:
|
||||
int nassets
|
||||
str column_name
|
||||
ndarray[dtype=uint32_t, ndim=1] raw_data
|
||||
ndarray[dtype=uint32_t, ndim=2] outbuf
|
||||
ndarray[dtype=uint8_t, ndim=2, cast=True] where_nan
|
||||
ndarray[dtype=float64_t, ndim=2] outbuf_as_float
|
||||
intp_t asset
|
||||
intp_t out_idx
|
||||
intp_t raw_idx
|
||||
intp_t first_row
|
||||
intp_t last_row
|
||||
intp_t offset
|
||||
list results = []
|
||||
|
||||
nassets = shape[1]
|
||||
if not nassets== len(first_rows) == len(last_rows) == len(offsets):
|
||||
raise ValueError("Incompatible index arrays.")
|
||||
|
||||
for column_name in columns:
|
||||
raw_data = table[column_name][:]
|
||||
outbuf = zeros(shape=shape, dtype=uint32)
|
||||
for asset in range(nassets):
|
||||
first_row = first_rows[asset]
|
||||
last_row = last_rows[asset]
|
||||
offset = offsets[asset]
|
||||
for out_idx, raw_idx in enumerate(range(first_row, last_row + 1)):
|
||||
outbuf[out_idx + offset, asset] = raw_data[raw_idx]
|
||||
|
||||
if column_name in {'open', 'high', 'low', 'close'}:
|
||||
where_nan = (outbuf == 0)
|
||||
outbuf_as_float = outbuf.astype(float64) * .001
|
||||
outbuf_as_float[where_nan] = NAN
|
||||
results.append(outbuf_as_float)
|
||||
else:
|
||||
results.append(outbuf)
|
||||
return results
|
||||
|
||||
@@ -27,7 +27,6 @@ from bcolz import (
|
||||
from click import progressbar
|
||||
from numpy import (
|
||||
array,
|
||||
array_equal,
|
||||
float64,
|
||||
floating,
|
||||
full,
|
||||
@@ -50,11 +49,11 @@ import sqlite3
|
||||
|
||||
|
||||
from zipline.data.ffc.base import FFCLoader
|
||||
from zipline.data.ffc.loaders._us_equity_pricing import (
|
||||
from zipline.data.ffc.loaders._equities import (
|
||||
_compute_row_slices,
|
||||
_read_bcolz_data,
|
||||
load_adjustments_from_sqlite,
|
||||
)
|
||||
from zipline.data.ffc.loaders._adjustments import load_adjustments_from_sqlite
|
||||
from zipline.lib.adjusted_array import (
|
||||
adjusted_array,
|
||||
)
|
||||
@@ -351,48 +350,17 @@ class BcolzDailyBarReader(object):
|
||||
for id_, offset in iteritems(table.attrs['calendar_offset'])
|
||||
}
|
||||
|
||||
def _slice_locs(self, start_date, end_date):
|
||||
try:
|
||||
start = self._calendar.get_loc(start_date)
|
||||
except KeyError:
|
||||
if start_date < self._calendar[0]:
|
||||
raise NoFurtherDataError(
|
||||
msg=(
|
||||
"FFC Query requesting data starting on {query_start}, "
|
||||
"but first known date is {calendar_start}"
|
||||
).format(
|
||||
query_start=str(start_date),
|
||||
calendar_start=str(self._calendar[0]),
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError("Query start %s not in calendar" % start_date)
|
||||
try:
|
||||
stop = self._calendar.get_loc(end_date)
|
||||
except:
|
||||
if end_date > self._calendar[-1]:
|
||||
raise NoFurtherDataError(
|
||||
msg=(
|
||||
"FFC Query requesting data up to {query_end}, "
|
||||
"but last known date is {calendar_end}"
|
||||
).format(
|
||||
query_end=end_date,
|
||||
calendar_end=self._calendar[-1],
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError("Query end %s not in calendar" % end_date)
|
||||
return start, stop
|
||||
|
||||
def _compute_slices(self, dates, assets):
|
||||
def _compute_slices(self, start_idx, end_idx, assets):
|
||||
"""
|
||||
Compute the raw row indices to load for each asset on a query for the
|
||||
given dates.
|
||||
given dates after applying a shift.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dates : pandas.DatetimeIndex
|
||||
Dates of the query on which we want to compute row indices.
|
||||
start_idx : int
|
||||
Index of first date for which we want data.
|
||||
end_idx : int
|
||||
Index of last date for which we want data.
|
||||
assets : pandas.Int64Index
|
||||
Assets for which we want to compute row indices
|
||||
|
||||
@@ -413,31 +381,29 @@ class BcolzDailyBarReader(object):
|
||||
of a query. Otherwise, offset[i] will be equal to the number of
|
||||
entries in `dates` for which the asset did not yet exist.
|
||||
"""
|
||||
start, stop = self._slice_locs(dates[0], dates[-1])
|
||||
|
||||
# Sanity check that the requested date range matches our calendar.
|
||||
# This could be removed in the future if it's materially affecting
|
||||
# performance.
|
||||
query_dates = self._calendar[start:stop + 1]
|
||||
if not array_equal(query_dates.values, dates.values):
|
||||
raise ValueError("Incompatible calendars!")
|
||||
|
||||
# The core implementation of the logic here is implemented in Cython
|
||||
# for efficiency.
|
||||
return _compute_row_slices(
|
||||
self._first_rows,
|
||||
self._last_rows,
|
||||
self._calendar_offsets,
|
||||
start,
|
||||
stop,
|
||||
start_idx,
|
||||
end_idx,
|
||||
assets,
|
||||
)
|
||||
|
||||
def load_raw_arrays(self, columns, dates, assets):
|
||||
first_rows, last_rows, offsets = self._compute_slices(dates, assets)
|
||||
def load_raw_arrays(self, columns, start_date, end_date, assets):
|
||||
# Assumes that the given dates are actually in calendar.
|
||||
start_idx = self._calendar.get_loc(start_date)
|
||||
end_idx = self._calendar.get_loc(end_date)
|
||||
first_rows, last_rows, offsets = self._compute_slices(
|
||||
start_idx,
|
||||
end_idx,
|
||||
assets,
|
||||
)
|
||||
return _read_bcolz_data(
|
||||
self._table,
|
||||
(len(dates), len(assets)),
|
||||
(end_idx - start_idx + 1, len(assets)),
|
||||
[column.name for column in columns],
|
||||
first_rows,
|
||||
last_rows,
|
||||
@@ -617,15 +583,28 @@ class USEquityPricingLoader(FFCLoader):
|
||||
|
||||
def __init__(self, raw_price_loader, adjustments_loader):
|
||||
self.raw_price_loader = raw_price_loader
|
||||
# HACK: Pull the calendar off our raw_price_loader so that we can
|
||||
# backshift dates.
|
||||
self._calendar = self.raw_price_loader._calendar
|
||||
self.adjustments_loader = adjustments_loader
|
||||
|
||||
def load_adjusted_array(self, columns, mask):
|
||||
dates, assets = mask.index, mask.columns
|
||||
def load_adjusted_array(self, columns, dates, assets, mask):
|
||||
# load_adjusted_array is called with dates on which the user's algo
|
||||
# will be shown data, which means we need to return the data that would
|
||||
# be known at the start of each date. We assume that the latest data
|
||||
# known on day N is the data from day (N - 1), so we shift all query
|
||||
# dates back by a day.
|
||||
start_date, end_date = _shift_dates(
|
||||
self._calendar, dates[0], dates[-1], shift=1,
|
||||
)
|
||||
|
||||
raw_arrays = self.raw_price_loader.load_raw_arrays(
|
||||
columns,
|
||||
dates,
|
||||
start_date,
|
||||
end_date,
|
||||
assets,
|
||||
)
|
||||
|
||||
adjustments = self.adjustments_loader.load_adjustments(
|
||||
columns,
|
||||
dates,
|
||||
@@ -633,6 +612,51 @@ class USEquityPricingLoader(FFCLoader):
|
||||
)
|
||||
|
||||
return [
|
||||
adjusted_array(raw_array, mask.values, col_adjustments)
|
||||
adjusted_array(raw_array, mask, col_adjustments)
|
||||
for raw_array, col_adjustments in zip(raw_arrays, adjustments)
|
||||
]
|
||||
|
||||
|
||||
def _shift_dates(dates, start_date, end_date, shift):
|
||||
try:
|
||||
start = dates.get_loc(start_date)
|
||||
except KeyError:
|
||||
if start_date < dates[0]:
|
||||
raise NoFurtherDataError(
|
||||
msg=(
|
||||
"Modeling Query requested data starting on {query_start}, "
|
||||
"but first known date is {calendar_start}"
|
||||
).format(
|
||||
query_start=str(start_date),
|
||||
calendar_start=str(dates[0]),
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError("Query start %s not in calendar" % start_date)
|
||||
|
||||
# Make sure that shifting doesn't push us out of the calendar.
|
||||
if start < shift:
|
||||
raise NoFurtherDataError(
|
||||
msg=(
|
||||
"Modeling Query requested data from {shift}"
|
||||
" days before {query_start}, but first known date is only "
|
||||
"{start} days earlier."
|
||||
).format(shift=shift, query_start=start_date, start=start),
|
||||
)
|
||||
|
||||
try:
|
||||
end = dates.get_loc(end_date)
|
||||
except KeyError:
|
||||
if end_date > dates[-1]:
|
||||
raise NoFurtherDataError(
|
||||
msg=(
|
||||
"Modeling Query requesting data up to {query_end}, "
|
||||
"but last known date is {calendar_end}"
|
||||
).format(
|
||||
query_end=end_date,
|
||||
calendar_end=dates[-1],
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError("Query end %s not in calendar" % end_date)
|
||||
return dates[start - shift], dates[end - shift]
|
||||
|
||||
@@ -48,17 +48,17 @@ class MultiColumnLoader(FFCLoader):
|
||||
def __init__(self, loaders):
|
||||
self._loaders = loaders
|
||||
|
||||
def load_adjusted_array(self, columns, mask):
|
||||
def load_adjusted_array(self, columns, dates, assets, mask):
|
||||
"""
|
||||
Load by delegating to sub-loaders.
|
||||
"""
|
||||
out = []
|
||||
for column in columns:
|
||||
for col in columns:
|
||||
try:
|
||||
loader = self._loaders[column]
|
||||
loader = self._loaders[col]
|
||||
except KeyError:
|
||||
raise ValueError("Couldn't find loader for %s" % column)
|
||||
out.extend(loader.load_adjusted_array([column], mask))
|
||||
raise ValueError("Couldn't find loader for %s" % col)
|
||||
out.extend(loader.load_adjusted_array([col], dates, assets, mask))
|
||||
return out
|
||||
|
||||
|
||||
|
||||
@@ -344,6 +344,16 @@ class WindowLengthNotSpecified(ZiplineError):
|
||||
)
|
||||
|
||||
|
||||
class DTypeNotSpecified(ZiplineError):
|
||||
"""
|
||||
Raised if a user attempts to construct a term without specifying dtype and
|
||||
that term does not have class-level default dtype.
|
||||
"""
|
||||
msg = (
|
||||
"{termname} requires a dtype, but no dtype was passed."
|
||||
)
|
||||
|
||||
|
||||
class BadPercentileBounds(ZiplineError):
|
||||
"""
|
||||
Raised by API functions accepting percentile bounds when the passed bounds
|
||||
|
||||
@@ -43,8 +43,6 @@ cpdef tuple get_adjustment_locs(DatetimeIndex_t dates_index,
|
||||
start_date_loc = dates_index.get_loc(start_date, method='bfill')
|
||||
|
||||
return (
|
||||
# start_date is allowed to be None, indicating "everything
|
||||
# before the end_date"
|
||||
start_date_loc,
|
||||
# Location of latest date on or before start_date.
|
||||
dates_index.get_loc(end_date, method='ffill'),
|
||||
@@ -135,6 +133,7 @@ cdef class Float64Adjustment:
|
||||
(other.first_row, other.last_row, other.col, other.value)
|
||||
)
|
||||
|
||||
|
||||
cdef class Float64Multiply(Float64Adjustment):
|
||||
"""
|
||||
An adjustment that multiplies by a scalar.
|
||||
|
||||
+123
-65
@@ -28,11 +28,13 @@ from pandas import (
|
||||
|
||||
from zipline.lib.adjusted_array import ensure_ndarray
|
||||
from zipline.errors import NoFurtherDataError
|
||||
from zipline.utils.pandas_utils import explode
|
||||
|
||||
from .classifier import Classifier
|
||||
from .factor import Factor
|
||||
from .filter import Filter
|
||||
from .graph import TermGraph
|
||||
from .term import AssetExists
|
||||
|
||||
|
||||
class FFCEngine(with_metaclass(ABCMeta)):
|
||||
@@ -49,17 +51,18 @@ class FFCEngine(with_metaclass(ABCMeta)):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
terms : dict
|
||||
Map from str -> zipline.modelling.term.Term.
|
||||
start_date : datetime
|
||||
The first date of the matrix.
|
||||
end_date : datetime
|
||||
The last date of the matrix.
|
||||
terms : dict[str -> zipline.modelling.term.Term]
|
||||
Dict mapping term names to instances. The supplied names are used
|
||||
as column names in our output frame.
|
||||
start_date : pd.Timestamp
|
||||
Start date of the computed matrix.
|
||||
end_date : pd.Timestamp
|
||||
End date of the computed matrix.
|
||||
|
||||
Returns
|
||||
-------
|
||||
matrix : pd.DataFrame
|
||||
A matrix of factors
|
||||
A matrix of computed results.
|
||||
"""
|
||||
raise NotImplementedError("factor_matrix")
|
||||
|
||||
@@ -97,6 +100,7 @@ class SimpleFFCEngine(object):
|
||||
'_loader',
|
||||
'_calendar',
|
||||
'_finder',
|
||||
'_root_mask_term',
|
||||
'__weakref__',
|
||||
]
|
||||
|
||||
@@ -104,6 +108,7 @@ class SimpleFFCEngine(object):
|
||||
self._loader = loader
|
||||
self._calendar = calendar
|
||||
self._finder = asset_finder
|
||||
self._root_mask_term = AssetExists()
|
||||
|
||||
def factor_matrix(self, terms, start_date, end_date):
|
||||
"""
|
||||
@@ -144,7 +149,7 @@ class SimpleFFCEngine(object):
|
||||
5. Stick the values computed in (4) into a DataFrame and return it.
|
||||
|
||||
Step 0 is performed by `zipline.modelling.graph.TermGraph`.
|
||||
Step 1 is performed in `self.build_lifetimes_matrix`.
|
||||
Step 1 is performed in `self._compute_root_mask`.
|
||||
Step 2 is performed in `self.compute_chunk`.
|
||||
Steps 3, 4, and 5 are performed in self._format_factor_matrix.
|
||||
|
||||
@@ -159,20 +164,18 @@ class SimpleFFCEngine(object):
|
||||
)
|
||||
|
||||
graph = TermGraph(terms)
|
||||
max_extra_rows = graph.max_extra_rows
|
||||
extra_rows = graph.extra_rows[self._root_mask_term]
|
||||
|
||||
lifetimes = self.build_lifetimes_matrix(
|
||||
start_date,
|
||||
end_date,
|
||||
max_extra_rows,
|
||||
root_mask = self._compute_root_mask(start_date, end_date, extra_rows)
|
||||
dates, assets, root_mask_values = explode(root_mask)
|
||||
raw_outputs = self.compute_chunk(
|
||||
graph,
|
||||
dates,
|
||||
assets,
|
||||
initial_workspace={self._root_mask_term: root_mask_values},
|
||||
)
|
||||
raw_outputs = self.compute_chunk(graph, lifetimes, {})
|
||||
|
||||
lifetimes_between_dates = lifetimes[max_extra_rows:]
|
||||
dates = lifetimes_between_dates.index.values
|
||||
assets = lifetimes_between_dates.columns.values
|
||||
|
||||
# We only need filters and factors to compute the final output matrix.
|
||||
# Collect the results that we'll actually show to the user.
|
||||
filters, factors = {}, {}
|
||||
for name, term in iteritems(terms):
|
||||
if isinstance(term, Filter):
|
||||
@@ -184,12 +187,14 @@ class SimpleFFCEngine(object):
|
||||
else:
|
||||
raise ValueError("Unknown term type: %s" % term)
|
||||
|
||||
# Treat base_mask as an implicit filter.
|
||||
# TODO: Is there a clean way to make this actually just be a filter?
|
||||
filters['base'] = lifetimes_between_dates.values
|
||||
return self._format_factor_matrix(dates, assets, filters, factors)
|
||||
# Add the root mask as an implicit filter, truncating off the extra
|
||||
# rows that we only needed to compute other terms.
|
||||
filters['base'] = root_mask_values[extra_rows:]
|
||||
out_dates = dates[extra_rows:]
|
||||
|
||||
def build_lifetimes_matrix(self, start_date, end_date, extra_rows):
|
||||
return self._format_factor_matrix(out_dates, assets, filters, factors)
|
||||
|
||||
def _compute_root_mask(self, start_date, end_date, extra_rows):
|
||||
"""
|
||||
Compute a lifetimes matrix from our AssetFinder, then drop columns that
|
||||
didn't exist at all during the query dates.
|
||||
@@ -201,9 +206,9 @@ class SimpleFFCEngine(object):
|
||||
end_date : pd.Timestamp
|
||||
End date for the matrix.
|
||||
extra_rows : int
|
||||
Number of rows prior to `start_date` to include.
|
||||
Number of extra rows to compute before `start_date`.
|
||||
Extra rows are needed by terms like moving averages that require a
|
||||
trailing window of data to compute.
|
||||
trailing window of data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -230,8 +235,10 @@ class SimpleFFCEngine(object):
|
||||
# Build lifetimes matrix reaching back to `extra_rows` days before
|
||||
# `start_date.`
|
||||
lifetimes = finder.lifetimes(
|
||||
calendar[start_idx - extra_rows:end_idx]
|
||||
calendar[start_idx - extra_rows:end_idx],
|
||||
include_start_date=False
|
||||
)
|
||||
|
||||
assert lifetimes.index[extra_rows] == start_date
|
||||
assert lifetimes.index[-1] == end_date
|
||||
if not lifetimes.columns.unique:
|
||||
@@ -244,36 +251,48 @@ class SimpleFFCEngine(object):
|
||||
existed = lifetimes.iloc[extra_rows:].any()
|
||||
return lifetimes.loc[:, existed]
|
||||
|
||||
def _inputs_for_term(self, term, workspace, extra_rows):
|
||||
def _mask_and_dates_for_term(self, term, workspace, graph, dates):
|
||||
"""
|
||||
Load mask and mask row labels for term.
|
||||
"""
|
||||
mask = term.mask
|
||||
offset = graph.extra_rows[mask] - graph.extra_rows[term]
|
||||
return workspace[mask][offset:], dates[offset:]
|
||||
|
||||
def _inputs_for_term(self, term, workspace, graph):
|
||||
"""
|
||||
Compute inputs for the given term.
|
||||
|
||||
This is mostly complicated by the fact that for each input we store
|
||||
as many rows as will be necessary to serve any term requiring that
|
||||
input. Thus if Factor A needs 5 extra rows of price, and Factor B
|
||||
needs 3 extra rows of price, we need to remove 2 leading rows from our
|
||||
stored prices before passing them to Factor B.
|
||||
This is mostly complicated by the fact that for each input we store as
|
||||
many rows as will be necessary to serve **any** computation requiring
|
||||
that input.
|
||||
"""
|
||||
term_extra_rows = term.extra_input_rows
|
||||
offsets = graph.offset
|
||||
if term.windowed:
|
||||
# If term is windowed, then all input data should be instances of
|
||||
# AdjustedArray.
|
||||
return [
|
||||
workspace[input_].traverse(
|
||||
term.window_length,
|
||||
offset=extra_rows[input_] - term_extra_rows
|
||||
)
|
||||
for input_ in term.inputs
|
||||
]
|
||||
else:
|
||||
return [
|
||||
ensure_ndarray(
|
||||
workspace[input_][
|
||||
extra_rows[input_] - term_extra_rows:
|
||||
],
|
||||
window_length=term.window_length,
|
||||
offset=offsets[term, input_]
|
||||
)
|
||||
for input_ in term.inputs
|
||||
]
|
||||
|
||||
def compute_chunk(self, graph, base_mask, initial_workspace):
|
||||
# If term is not windowed, input_data may be an AdjustedArray or
|
||||
# np.ndarray. Coerce the former to the latter.
|
||||
out = []
|
||||
for input_ in term.inputs:
|
||||
input_data = ensure_ndarray(workspace[input_])
|
||||
offset = offsets[term, input_]
|
||||
# OPTIMIZATION: Don't make a copy by doing input_data[0:] if
|
||||
# offset is zero.
|
||||
if offset:
|
||||
input_data = input_data[offset:]
|
||||
out.append(input_data)
|
||||
return out
|
||||
|
||||
def compute_chunk(self, graph, dates, assets, initial_workspace):
|
||||
"""
|
||||
Compute the FFC terms in the graph for the requested start and end
|
||||
dates.
|
||||
@@ -281,53 +300,64 @@ class SimpleFFCEngine(object):
|
||||
Parameters
|
||||
----------
|
||||
graph : zipline.modelling.graph.TermGraph
|
||||
dates : pd.DatetimeIndex
|
||||
Row labels for our root mask.
|
||||
assets : pd.Int64Index
|
||||
Column labels for our root mask.
|
||||
initial_workspace : dict
|
||||
Map from term -> output.
|
||||
Must contain at least entry for `self._root_mask_term` whose shape
|
||||
is `(len(dates), len(assets))`, but may contain additional
|
||||
pre-computed terms for testing or optimization purposes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
results : dict
|
||||
Dictionary mapping requested results to outputs.
|
||||
"""
|
||||
self._validate_compute_chunk_params(dates, assets, initial_workspace)
|
||||
loader = self._loader
|
||||
|
||||
extra_rows = graph.extra_rows
|
||||
max_extra_rows = graph.max_extra_rows
|
||||
|
||||
workspace = {}
|
||||
if initial_workspace is not None:
|
||||
workspace.update(initial_workspace)
|
||||
# Copy the supplied initial workspace so we don't mutate it in place.
|
||||
workspace = initial_workspace.copy()
|
||||
|
||||
for term in graph.ordered():
|
||||
# Subclasses are allowed to pre-populate computed values for terms,
|
||||
# and in the future we may pre-compute atomic terms coming from the
|
||||
# same dataset. In both cases, it's possible that we already have
|
||||
# an entry for this term.
|
||||
# `term` may have been supplied in `initial_workspace`, and in the
|
||||
# future we may pre-compute atomic terms coming from the same
|
||||
# dataset. In either case, we will already have an entry for this
|
||||
# term, which we shouldn't re-compute.
|
||||
if term in workspace:
|
||||
continue
|
||||
|
||||
base_mask_for_term = base_mask.iloc[
|
||||
max_extra_rows - extra_rows[term]:
|
||||
]
|
||||
# Asset labels are always the same, but date labels vary by how
|
||||
# many extra rows are needed.
|
||||
mask, mask_dates = self._mask_and_dates_for_term(
|
||||
term, workspace, graph, dates
|
||||
)
|
||||
if term.atomic:
|
||||
# FUTURE OPTIMIZATION: Scan the resolution order for terms in
|
||||
# the same dataset and load them here as well.
|
||||
to_load = [term]
|
||||
loaded = loader.load_adjusted_array(
|
||||
to_load,
|
||||
base_mask_for_term,
|
||||
to_load, mask_dates, assets, mask,
|
||||
)
|
||||
assert len(to_load) == len(loaded)
|
||||
for loaded_term, adj_array in zip_longest(to_load, loaded):
|
||||
workspace[loaded_term] = adj_array
|
||||
else:
|
||||
workspace[term] = term._compute(
|
||||
self._inputs_for_term(term, workspace, extra_rows),
|
||||
base_mask_for_term,
|
||||
self._inputs_for_term(term, workspace, graph),
|
||||
mask_dates,
|
||||
assets,
|
||||
mask,
|
||||
)
|
||||
assert(workspace[term].shape == base_mask_for_term.shape)
|
||||
assert(workspace[term].shape == mask.shape)
|
||||
|
||||
out = {}
|
||||
graph_extra_rows = graph.extra_rows
|
||||
for name, term in iteritems(graph.outputs):
|
||||
# Truncate off extra rows from outputs.
|
||||
out[name] = workspace[term][extra_rows[term]:]
|
||||
out[name] = workspace[term][graph_extra_rows[term]:]
|
||||
return out
|
||||
|
||||
def _format_factor_matrix(self, dates, assets, filters, factors):
|
||||
@@ -423,3 +453,31 @@ class SimpleFFCEngine(object):
|
||||
],
|
||||
)
|
||||
).tz_localize('UTC', level=0)
|
||||
|
||||
def _validate_compute_chunk_params(self, dates, assets, initial_workspace):
|
||||
"""
|
||||
Verify that the values passed to compute_chunk are well-formed.
|
||||
"""
|
||||
root = self._root_mask_term
|
||||
clsname = type(self).__name__
|
||||
# Writing this out explicitly so this errors in testing if we change
|
||||
# the name without updating this line.
|
||||
compute_chunk_name = self.compute_chunk.__name__
|
||||
if root not in initial_workspace:
|
||||
raise AssertionError(
|
||||
"root_mask values not supplied to {cls}.{method}".format(
|
||||
cls=clsname,
|
||||
method=compute_chunk_name,
|
||||
)
|
||||
)
|
||||
|
||||
shape = initial_workspace[root].shape
|
||||
implied_shape = len(dates), len(assets)
|
||||
if shape != implied_shape:
|
||||
raise AssertionError(
|
||||
"root_mask shape is {shape}, but received dates/assets "
|
||||
"imply that shape should be {implied}".format(
|
||||
shape=shape,
|
||||
implied=implied_shape,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -12,7 +12,7 @@ from numpy import (
|
||||
)
|
||||
from six import integer_types
|
||||
|
||||
from zipline.modelling.term import Term
|
||||
from zipline.modelling.term import Term, NotSpecified
|
||||
|
||||
_VARIABLE_NAME_RE = re.compile("^(x_)([0-9]+)$")
|
||||
|
||||
@@ -180,7 +180,7 @@ class NumericalExpression(Term):
|
||||
# factor(int64) + factor(int64) + 2.5.
|
||||
# The real fix for this is probably for the calling context to specify
|
||||
# dtypes.
|
||||
if cls.dtype is not None:
|
||||
if cls.dtype is not NotSpecified:
|
||||
dtype = cls.dtype
|
||||
else:
|
||||
dtype = find_common_type(
|
||||
@@ -228,7 +228,7 @@ class NumericalExpression(Term):
|
||||
)
|
||||
return super(NumericalExpression, self)._validate()
|
||||
|
||||
def _compute(self, arrays, mask):
|
||||
def _compute(self, arrays, dates, assets, mask):
|
||||
"""
|
||||
Compute our stored expression string with numexpr.
|
||||
"""
|
||||
|
||||
@@ -16,6 +16,7 @@ from zipline.errors import (
|
||||
from zipline.lib.rank import rankdata_2d_ordinal
|
||||
from zipline.modelling.term import (
|
||||
CustomTermMixin,
|
||||
NotSpecified,
|
||||
RequiredWindowLengthMixin,
|
||||
SingleInputMixin,
|
||||
Term,
|
||||
@@ -184,6 +185,8 @@ class Factor(Term):
|
||||
A transformation yielding a timeseries of scalar values associated with an
|
||||
Asset.
|
||||
"""
|
||||
dtype = float64
|
||||
|
||||
# Dynamically add functions for creating NumExprFactor/NumExprFilter
|
||||
# instances.
|
||||
clsdict = locals()
|
||||
@@ -219,7 +222,7 @@ class Factor(Term):
|
||||
|
||||
eq = binary_operator('==')
|
||||
|
||||
def rank(self, method='ordinal', ascending=True):
|
||||
def rank(self, method='ordinal', ascending=True, mask=NotSpecified):
|
||||
"""
|
||||
Construct a new Factor representing the sorted rank of each column
|
||||
within each row.
|
||||
@@ -256,9 +259,9 @@ class Factor(Term):
|
||||
zipline.lib.rank
|
||||
zipline.modelling.factor.Rank
|
||||
"""
|
||||
return Rank(self if ascending else -self, method=method)
|
||||
return Rank(self if ascending else -self, method=method, mask=mask)
|
||||
|
||||
def top(self, N):
|
||||
def top(self, N, mask=NotSpecified):
|
||||
"""
|
||||
Construct a Filter matching the top N asset values of self each day.
|
||||
|
||||
@@ -271,9 +274,9 @@ class Factor(Term):
|
||||
-------
|
||||
filter : zipline.modelling.filter.Filter
|
||||
"""
|
||||
return self.rank(ascending=False) <= N
|
||||
return self.rank(ascending=False, mask=mask) <= N
|
||||
|
||||
def bottom(self, N):
|
||||
def bottom(self, N, mask=NotSpecified):
|
||||
"""
|
||||
Construct a Filter matching the bottom N asset values of self each day.
|
||||
|
||||
@@ -286,9 +289,12 @@ class Factor(Term):
|
||||
-------
|
||||
filter : zipline.modelling.filter.Filter
|
||||
"""
|
||||
return self.rank(ascending=True) <= N
|
||||
return self.rank(ascending=True, mask=mask) <= N
|
||||
|
||||
def percentile_between(self, min_percentile, max_percentile):
|
||||
def percentile_between(self,
|
||||
min_percentile,
|
||||
max_percentile,
|
||||
mask=NotSpecified):
|
||||
"""
|
||||
Construct a new Filter representing entries from the output of this
|
||||
Factor that fall within the percentile range defined by min_percentile
|
||||
@@ -312,6 +318,7 @@ class Factor(Term):
|
||||
self,
|
||||
min_percentile=min_percentile,
|
||||
max_percentile=max_percentile,
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
|
||||
@@ -359,15 +366,15 @@ class Rank(SingleInputMixin, Factor):
|
||||
Most users should call Factor.rank rather than directly construct an
|
||||
instance of this class.
|
||||
"""
|
||||
dtype = float64
|
||||
window_length = 0
|
||||
domain = None
|
||||
dtype = float64
|
||||
|
||||
def __new__(cls, factor, method):
|
||||
def __new__(cls, factor, method, mask):
|
||||
return super(Rank, cls).__new__(
|
||||
cls,
|
||||
inputs=(factor,),
|
||||
method=method,
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
def _init(self, method, *args, **kwargs):
|
||||
@@ -392,35 +399,35 @@ class Rank(SingleInputMixin, Factor):
|
||||
)
|
||||
return super(Rank, self)._validate()
|
||||
|
||||
def _compute(self, arrays, mask):
|
||||
def _compute(self, arrays, dates, assets, mask):
|
||||
"""
|
||||
For each row in the input, compute a like-shaped array of per-row
|
||||
ranks.
|
||||
"""
|
||||
# OPTIMIZATION: Fast path the default value with our own specialized
|
||||
# implementation.
|
||||
inv_mask = ~mask
|
||||
data = arrays[0].copy()
|
||||
data[inv_mask] = nan
|
||||
# OPTIMIZATION: Fast path the default case with our own specialized
|
||||
# Cython implementation.
|
||||
if self._method == 'ordinal':
|
||||
result = rankdata_2d_ordinal(arrays[0])
|
||||
result = rankdata_2d_ordinal(data)
|
||||
else:
|
||||
# FUTURE OPTIMIZATION:
|
||||
# Write a less general "apply to rows" method that doesn't do all
|
||||
# the extra work that apply_along_axis does.
|
||||
result = apply_along_axis(
|
||||
rankdata,
|
||||
1,
|
||||
arrays[0],
|
||||
method=self._method,
|
||||
)
|
||||
# rankdata will sort nan values into last place, but we want our
|
||||
# nans to propagate, so explicitly re-apply
|
||||
result[~mask.values] = nan
|
||||
result = apply_along_axis(rankdata, 1, data, method=self._method)
|
||||
|
||||
# rankdata will sort nan values into last place, but we want our
|
||||
# nans to propagate, so explicitly re-apply.
|
||||
result[inv_mask] = nan
|
||||
return result
|
||||
|
||||
def __repr__(self):
|
||||
return "{type}({input_}, method='{method}')".format(
|
||||
return "{type}({input_}, method='{method}', mask={mask})".format(
|
||||
type=type(self).__name__,
|
||||
input_=self.inputs[0],
|
||||
method=self._method,
|
||||
mask=self.mask,
|
||||
)
|
||||
|
||||
|
||||
@@ -433,7 +440,6 @@ class CustomFactor(RequiredWindowLengthMixin, CustomTermMixin, Factor):
|
||||
|
||||
We currently only support CustomFactors of type float64.
|
||||
"""
|
||||
dtype = float64
|
||||
ctx = nullctx()
|
||||
|
||||
def _validate(self):
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
Factor that produces the most most recently-known value of Column.
|
||||
"""
|
||||
from .factor import CustomFactor
|
||||
from ..term import SingleInputMixin
|
||||
|
||||
|
||||
class Latest(SingleInputMixin, CustomFactor):
|
||||
"""
|
||||
Factor producing the most recently-known value of `inputs[0]` on each day.
|
||||
"""
|
||||
window_length = 1
|
||||
|
||||
def compute(self, today, assets, out, data):
|
||||
out[:] = data[-1]
|
||||
+10
-92
@@ -85,7 +85,6 @@ class Filter(Term):
|
||||
"""
|
||||
A boolean predicate on a universe of Assets.
|
||||
"""
|
||||
domain = None
|
||||
dtype = bool_
|
||||
|
||||
clsdict = locals()
|
||||
@@ -96,41 +95,22 @@ class Filter(Term):
|
||||
}
|
||||
)
|
||||
|
||||
def then(self, other):
|
||||
"""
|
||||
Create a new filter by computing `self`, then computing `other` on the
|
||||
data that survived the first filter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
other : zipline.modelling.filter.Filter
|
||||
The Filter to apply next.
|
||||
|
||||
Returns
|
||||
-------
|
||||
filter : zipline.modelling.filter.SequencedFilter
|
||||
A filter which will compute `self` and then `other`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
zipline.modelling.filter.SequencedFilter
|
||||
"""
|
||||
return SequencedFilter(self, other)
|
||||
|
||||
|
||||
class NumExprFilter(NumericalExpression, Filter):
|
||||
"""
|
||||
A Filter computed from a numexpr expression.
|
||||
"""
|
||||
|
||||
def _compute(self, arrays, mask):
|
||||
def _compute(self, arrays, dates, assets, mask):
|
||||
"""
|
||||
Compute our result with numexpr, then apply `mask`.
|
||||
Compute our result with numexpr, then re-apply `mask`.
|
||||
"""
|
||||
return super(NumExprFilter, self)._compute(
|
||||
arrays,
|
||||
dates,
|
||||
assets,
|
||||
mask,
|
||||
) & mask.values
|
||||
) & mask
|
||||
|
||||
|
||||
class PercentileFilter(SingleInputMixin, Filter):
|
||||
@@ -148,10 +128,11 @@ class PercentileFilter(SingleInputMixin, Filter):
|
||||
"""
|
||||
window_length = 0
|
||||
|
||||
def __new__(cls, factor, min_percentile, max_percentile):
|
||||
def __new__(cls, factor, min_percentile, max_percentile, mask):
|
||||
return super(PercentileFilter, cls).__new__(
|
||||
cls,
|
||||
inputs=(factor,),
|
||||
mask=mask,
|
||||
min_percentile=min_percentile,
|
||||
max_percentile=max_percentile,
|
||||
)
|
||||
@@ -180,15 +161,15 @@ class PercentileFilter(SingleInputMixin, Filter):
|
||||
)
|
||||
return super(PercentileFilter, self)._validate()
|
||||
|
||||
def _compute(self, arrays, mask):
|
||||
def _compute(self, arrays, dates, assets, mask):
|
||||
"""
|
||||
For each row in the input, compute a mask of all values falling between
|
||||
the given percentiles.
|
||||
"""
|
||||
# TODO: Review whether there's a better way of handling small numbers
|
||||
# of columns.
|
||||
data = arrays[0].astype(float64)
|
||||
data[~mask.values] = nan
|
||||
data = arrays[0].copy().astype(float64)
|
||||
data[~mask] = nan
|
||||
|
||||
# FIXME: np.nanpercentile **should** support computing multiple bounds
|
||||
# at once, but there's a bug in the logic for multiple bounds in numpy
|
||||
@@ -207,66 +188,3 @@ class PercentileFilter(SingleInputMixin, Filter):
|
||||
keepdims=True,
|
||||
)
|
||||
return (lower_bounds <= data) & (data <= upper_bounds)
|
||||
|
||||
|
||||
class SequencedFilter(Filter):
|
||||
"""
|
||||
Term representing sequenced computation of two Filters.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
first : zipline.modelling.filter.Filter
|
||||
The first filter to compute.
|
||||
second : zipline.modelling.filter.Filter
|
||||
The second filter to compute.
|
||||
|
||||
Notes
|
||||
-----
|
||||
In general, users should rarely have to construct SequencedFilter instances
|
||||
directly. Instead, prefer construction via `Filter.then`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
Filter.then
|
||||
"""
|
||||
window_length = 0
|
||||
|
||||
def __new__(cls, first, then):
|
||||
return super(SequencedFilter, cls).__new__(
|
||||
cls,
|
||||
inputs=concat_tuples((first,), then.inputs),
|
||||
then=then,
|
||||
)
|
||||
|
||||
def _init(self, then, *args, **kwargs):
|
||||
self._then = then
|
||||
return super(SequencedFilter, self)._init(*args, **kwargs)
|
||||
|
||||
def _validate(self):
|
||||
"""
|
||||
Ensure that we're actually sequencing filters.
|
||||
"""
|
||||
first, then = self.inputs[0], self._then
|
||||
if not isinstance(first, Filter):
|
||||
raise TypeError("Expected Filter, got %s" % type(first).__name__)
|
||||
if not isinstance(then, Filter):
|
||||
raise TypeError("Expected Filter, got %s" % type(then).__name__)
|
||||
return super(SequencedFilter, self)._validate()
|
||||
|
||||
@classmethod
|
||||
def static_identity(cls, then, *args, **kwargs):
|
||||
return (
|
||||
super(SequencedFilter, cls).static_identity(*args, **kwargs),
|
||||
then,
|
||||
)
|
||||
|
||||
def _compute(self, arrays, mask):
|
||||
"""
|
||||
Call our second filter on its inputs, masking out any inputs rejected
|
||||
by our first filter.
|
||||
"""
|
||||
first_result, then_inputs = arrays[0], arrays[1:]
|
||||
return self._then._compute(
|
||||
then_inputs,
|
||||
mask & first_result,
|
||||
)
|
||||
|
||||
+127
-36
@@ -6,6 +6,7 @@ from networkx import (
|
||||
topological_sort,
|
||||
)
|
||||
from six import itervalues, iteritems
|
||||
from zipline.utils.memoize import lazyval
|
||||
|
||||
|
||||
class CyclicDependency(Exception):
|
||||
@@ -14,7 +15,7 @@ class CyclicDependency(Exception):
|
||||
|
||||
class TermGraph(DiGraph):
|
||||
"""
|
||||
Graph represention of FFC Term dependencies
|
||||
Graph represention of FFC Term dependencies.
|
||||
|
||||
Each node in the graph has an `extra_rows` attribute, indicating how many,
|
||||
if any, extra rows we should compute for the node. Extra rows are most
|
||||
@@ -25,13 +26,13 @@ class TermGraph(DiGraph):
|
||||
Parameters
|
||||
----------
|
||||
terms : dict
|
||||
A dict mapping names to terms.
|
||||
A dict mapping names to final output terms.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
outputs
|
||||
offset
|
||||
extra_rows
|
||||
max_extra_rows
|
||||
|
||||
Methods
|
||||
-------
|
||||
@@ -47,12 +48,106 @@ class TermGraph(DiGraph):
|
||||
assert not parents
|
||||
|
||||
self._outputs = terms
|
||||
self._extra_rows = {
|
||||
self._ordered = topological_sort(self)
|
||||
|
||||
@lazyval
|
||||
def offset(self):
|
||||
"""
|
||||
For all pairs (term, input) such that `input` is an input to `term`,
|
||||
compute a mapping:
|
||||
|
||||
(term, input) -> offset(term, input)
|
||||
|
||||
where `offset(term, input)` is defined as
|
||||
|
||||
Max number of extra rows needed by any term depending on `input`
|
||||
minus
|
||||
Number of extra rows needed by `term`.
|
||||
|
||||
Example
|
||||
-------
|
||||
|
||||
Factor A needs 5 extra rows of USEquityPricing.close, and Factor B
|
||||
needs 3 extra rows of the same. Factor A also requires 5 extra rows of
|
||||
USEquityPricing.high, which no other Factor users.
|
||||
|
||||
We load 5 extra rows of both `price` and `high` to ensure we can
|
||||
service Factor A, and the following offsets get computed:
|
||||
|
||||
self.offset[Factor A, USEquityPricing.close] == 0
|
||||
self.offset[Factor A, USEquityPricing.high] == 0
|
||||
self.offset[Factor B, USEquityPricing.close] == 2
|
||||
self.offset[Factor B, USEquityPricing.high] raises KeyError.
|
||||
|
||||
Notes
|
||||
-----
|
||||
`offset(term, input) >= 0` for all valid pairs, since `input` must be
|
||||
an input to `term` if the pair appears in the mapping.
|
||||
|
||||
This value is useful because we load enough rows of each input to serve
|
||||
all possible dependencies. However, for any given dependency, we only
|
||||
want to compute using the actual number of required extra rows for that
|
||||
dependency. We can do so by truncating off the first `offset` rows of
|
||||
the loaded data for `input`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
zipline.modelling.graph.TermGraph.offset
|
||||
zipline.modelling.engine.SimpleFFCEngine._inputs_for_term
|
||||
zipline.modelling.engine.SimpleFFCEngine._mask_for_term
|
||||
"""
|
||||
out = {}
|
||||
for term in self:
|
||||
extra_input_rows = term.extra_input_rows
|
||||
for input_ in term.inputs:
|
||||
out[term, input_] = self.extra_rows[input_] - extra_input_rows
|
||||
mask = term.mask
|
||||
if term.mask is not None:
|
||||
out[term, mask] = self.extra_rows[mask] - extra_input_rows
|
||||
return out
|
||||
|
||||
@lazyval
|
||||
def extra_rows(self):
|
||||
"""
|
||||
A dict mapping `term` -> `# of extra rows to load/compute of `term`.
|
||||
|
||||
This is always the maximum number of extra **input** rows required by
|
||||
any Filter/Factor for which `term` is an input.
|
||||
|
||||
Notes
|
||||
----
|
||||
This value depends on the other terms in the graph that require `term`
|
||||
**as an input**. This is not to be confused with
|
||||
`term.extra_input_rows`, which is how many extra rows of `term`'s
|
||||
inputs we need to load, and which is determined entirely by `Term`
|
||||
itself.
|
||||
|
||||
Example
|
||||
-------
|
||||
Our graph contains the following terms:
|
||||
|
||||
A = SimpleMovingAverage([USEquityPricing.high], window_length=5)
|
||||
B = SimpleMovingAverage([USEquityPricing.high], window_length=10)
|
||||
C = SimpleMovingAverage([USEquityPricing.low], window_length=8)
|
||||
|
||||
To compute N rows of A, we need N + 4 extra rows of `high`.
|
||||
To compute N rows of B, we need N + 9 extra rows of `high`.
|
||||
To compute N rows of C, we need N + 7 extra rows of `low`.
|
||||
|
||||
We store the following extra_row requirements:
|
||||
|
||||
self.extra_rows[high] = 9 # Ensures that we can service B.
|
||||
self.extra_rows[low] = 7
|
||||
|
||||
See Also
|
||||
--------
|
||||
zipline.modelling.graph.TermGraph.offset
|
||||
zipline.modelling.term.Term.extra_input_rows
|
||||
"""
|
||||
return {
|
||||
term: attrs['extra_rows']
|
||||
for term, attrs in iteritems(self.node)
|
||||
}
|
||||
self._max_extra_rows = max(itervalues(self._extra_rows))
|
||||
self._ordered = topological_sort(self)
|
||||
|
||||
@property
|
||||
def outputs(self):
|
||||
@@ -61,20 +156,6 @@ class TermGraph(DiGraph):
|
||||
"""
|
||||
return self._outputs
|
||||
|
||||
@property
|
||||
def extra_rows(self):
|
||||
"""
|
||||
Dict mapping term -> number of extra rows to compute for term.
|
||||
"""
|
||||
return self._extra_rows
|
||||
|
||||
@property
|
||||
def max_extra_rows(self):
|
||||
"""
|
||||
Maximum number of extra rows required to compute any term in the graph.
|
||||
"""
|
||||
return self._max_extra_rows
|
||||
|
||||
def ordered(self):
|
||||
"""
|
||||
Return a topologically-sorted iterator over the terms in `self`.
|
||||
@@ -92,25 +173,35 @@ class TermGraph(DiGraph):
|
||||
raise CyclicDependency(term)
|
||||
parents.add(term)
|
||||
|
||||
try:
|
||||
existing = self.node[term]
|
||||
except KeyError:
|
||||
# `term` is not yet in the graph: add it with the specified number
|
||||
# of extra rows.
|
||||
self.add_node(term, extra_rows=extra_rows)
|
||||
else:
|
||||
# `term` is already in the graph because we've been traversed by
|
||||
# another parent. Ensure that we have enough extra rows to satisfy
|
||||
# all of our parents.
|
||||
existing['extra_rows'] = max(extra_rows, existing['extra_rows'])
|
||||
# Idempotent if term is already in the graph.
|
||||
self.add_node(term)
|
||||
|
||||
extra_rows_for_subterms = extra_rows + term.extra_input_rows
|
||||
for subterm in term.inputs:
|
||||
# Make sure we're going to compute at least `extra_rows` of `term`.
|
||||
self._ensure_extra_rows(term, extra_rows)
|
||||
|
||||
# Number of extra rows we need to compute for this term's dependencies.
|
||||
dependency_extra_rows = extra_rows + term.extra_input_rows
|
||||
|
||||
# Recursively add dependencies.
|
||||
for dependency in term.inputs:
|
||||
self._add_to_graph(
|
||||
subterm,
|
||||
dependency,
|
||||
parents,
|
||||
extra_rows=extra_rows_for_subterms
|
||||
extra_rows=dependency_extra_rows,
|
||||
)
|
||||
self.add_edge(subterm, term)
|
||||
self.add_edge(dependency, term)
|
||||
|
||||
# Add term's mask, which is really just a specially-enumerated input.
|
||||
mask = term.mask
|
||||
if mask is not None:
|
||||
self._add_to_graph(mask, parents, extra_rows=dependency_extra_rows)
|
||||
self.add_edge(mask, term)
|
||||
|
||||
parents.remove(term)
|
||||
|
||||
def _ensure_extra_rows(self, term, N):
|
||||
"""
|
||||
Ensure that we're going to compute at least N extra rows of `term`.
|
||||
"""
|
||||
attrs = self.node[term]
|
||||
attrs['extra_rows'] = max(N, attrs.get('extra_rows', 0))
|
||||
|
||||
+96
-29
@@ -1,15 +1,12 @@
|
||||
"""
|
||||
Base class for Filters, Factors and Classifiers
|
||||
"""
|
||||
from numpy import (
|
||||
empty,
|
||||
float64,
|
||||
full,
|
||||
nan,
|
||||
)
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
from numpy import bool_, full, nan
|
||||
|
||||
from zipline.errors import (
|
||||
DTypeNotSpecified,
|
||||
InputTermNotAtomic,
|
||||
TermInputsNotSpecified,
|
||||
WindowLengthNotPositive,
|
||||
@@ -18,25 +15,48 @@ from zipline.errors import (
|
||||
from zipline.utils.memoize import lazyval
|
||||
|
||||
|
||||
NotSpecified = (object(),)
|
||||
@object.__new__ # bind a single instance to the name 'NotSpecified'
|
||||
class NotSpecified(object):
|
||||
"""
|
||||
Singleton sentinel value used for Term defaults.
|
||||
"""
|
||||
__slots__ = ('__weakref__',)
|
||||
|
||||
def __new__(cls):
|
||||
raise TypeError("Can't construct new instances of NotSpecified")
|
||||
|
||||
def __repr__(self):
|
||||
return type(self).__name__
|
||||
|
||||
def __reduce__(self):
|
||||
return type(self).__name__
|
||||
|
||||
def __deepcopy__(self, _memo):
|
||||
return self
|
||||
|
||||
def __copy__(self):
|
||||
return self
|
||||
|
||||
|
||||
class Term(object):
|
||||
"""
|
||||
Base class for terms in an FFC API compute graph.
|
||||
"""
|
||||
# These are NotSpecified because a subclass is required to provide them.
|
||||
inputs = NotSpecified
|
||||
window_length = NotSpecified
|
||||
domain = None
|
||||
dtype = float64
|
||||
dtype = NotSpecified
|
||||
mask = NotSpecified
|
||||
domain = NotSpecified
|
||||
|
||||
_term_cache = WeakValueDictionary()
|
||||
|
||||
def __new__(cls,
|
||||
inputs=None,
|
||||
window_length=None,
|
||||
domain=None,
|
||||
dtype=None,
|
||||
inputs=NotSpecified,
|
||||
mask=NotSpecified,
|
||||
window_length=NotSpecified,
|
||||
domain=NotSpecified,
|
||||
dtype=NotSpecified,
|
||||
*args,
|
||||
**kwargs):
|
||||
"""
|
||||
@@ -49,22 +69,35 @@ class Term(object):
|
||||
Caching previously-constructed Terms is **sane** because terms and
|
||||
their inputs are both conceptually immutable.
|
||||
"""
|
||||
if inputs is None:
|
||||
inputs = tuple(cls.inputs)
|
||||
else:
|
||||
# Class-level attributes can be used to provide defaults for Term
|
||||
# subclasses.
|
||||
|
||||
if inputs is NotSpecified:
|
||||
inputs = cls.inputs
|
||||
# Having inputs = NotSpecified is an error, but we handle it later
|
||||
# in self._validate rather than here.
|
||||
if inputs is not NotSpecified:
|
||||
# Allow users to specify lists as class-level defaults, but
|
||||
# normalize to a tuple so that inputs is hashable.
|
||||
inputs = tuple(inputs)
|
||||
|
||||
if window_length is None:
|
||||
if mask is NotSpecified:
|
||||
mask = cls.mask
|
||||
if mask is NotSpecified:
|
||||
mask = AssetExists()
|
||||
|
||||
if window_length is NotSpecified:
|
||||
window_length = cls.window_length
|
||||
|
||||
if domain is None:
|
||||
if domain is NotSpecified:
|
||||
domain = cls.domain
|
||||
|
||||
if dtype is None:
|
||||
if dtype is NotSpecified:
|
||||
dtype = cls.dtype
|
||||
|
||||
identity = cls.static_identity(
|
||||
inputs=inputs,
|
||||
mask=mask,
|
||||
window_length=window_length,
|
||||
domain=domain,
|
||||
dtype=dtype,
|
||||
@@ -77,6 +110,7 @@ class Term(object):
|
||||
new_instance = cls._term_cache[identity] = \
|
||||
super(Term, cls).__new__(cls)._init(
|
||||
inputs=inputs,
|
||||
mask=mask,
|
||||
window_length=window_length,
|
||||
domain=domain,
|
||||
dtype=dtype,
|
||||
@@ -100,8 +134,9 @@ class Term(object):
|
||||
"""
|
||||
pass
|
||||
|
||||
def _init(self, inputs, window_length, domain, dtype):
|
||||
def _init(self, inputs, mask, window_length, domain, dtype):
|
||||
self.inputs = inputs
|
||||
self.mask = mask
|
||||
self.window_length = window_length
|
||||
self.domain = domain
|
||||
self.dtype = dtype
|
||||
@@ -110,7 +145,7 @@ class Term(object):
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def static_identity(cls, inputs, window_length, domain, dtype):
|
||||
def static_identity(cls, inputs, mask, window_length, domain, dtype):
|
||||
"""
|
||||
Return the identity of the Term that would be constructed from the
|
||||
given arguments.
|
||||
@@ -122,7 +157,7 @@ class Term(object):
|
||||
This is a classmethod so that it can be called from Term.__new__ to
|
||||
determine whether to produce a new instance.
|
||||
"""
|
||||
return (cls, inputs, window_length, domain, dtype)
|
||||
return (cls, inputs, mask, window_length, domain, dtype)
|
||||
|
||||
def _validate(self):
|
||||
"""
|
||||
@@ -133,6 +168,11 @@ class Term(object):
|
||||
raise TermInputsNotSpecified(termname=type(self).__name__)
|
||||
if self.window_length is NotSpecified:
|
||||
raise WindowLengthNotSpecified(termname=type(self).__name__)
|
||||
if self.dtype is NotSpecified:
|
||||
raise DTypeNotSpecified(termname=type(self).__name__)
|
||||
if self.mask is NotSpecified and not self.atomic:
|
||||
# This isn't user error, this is a bug in our code.
|
||||
raise AssertionError("{term} has no mask".format(term=self))
|
||||
|
||||
if self.window_length:
|
||||
for child in self.inputs:
|
||||
@@ -172,7 +212,7 @@ class Term(object):
|
||||
"""
|
||||
return max(0, self.window_length - 1)
|
||||
|
||||
def _compute(self, inputs, mask):
|
||||
def _compute(self, inputs, dates, assets, mask):
|
||||
"""
|
||||
Subclasses should implement this to perform actual computation.
|
||||
|
||||
@@ -188,6 +228,7 @@ class Term(object):
|
||||
type=type(self).__name__,
|
||||
inputs=self.inputs,
|
||||
window_length=self.window_length,
|
||||
mask=self.mask,
|
||||
)
|
||||
|
||||
|
||||
@@ -209,9 +250,9 @@ class SingleInputMixin(object):
|
||||
|
||||
class RequiredWindowLengthMixin(object):
|
||||
def _validate(self):
|
||||
if self.windowed or self.window_length is NotSpecified:
|
||||
return super(RequiredWindowLengthMixin, self)._validate()
|
||||
raise WindowLengthNotPositive(window_length=self.window_length)
|
||||
if not self.windowed:
|
||||
raise WindowLengthNotPositive(window_length=self.window_length)
|
||||
return super(RequiredWindowLengthMixin, self)._validate()
|
||||
|
||||
|
||||
class CustomTermMixin(object):
|
||||
@@ -230,14 +271,13 @@ class CustomTermMixin(object):
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def _compute(self, windows, mask):
|
||||
def _compute(self, windows, dates, assets, mask):
|
||||
"""
|
||||
Call the user's `compute` function on each window with a pre-built
|
||||
output array.
|
||||
"""
|
||||
# TODO: Make mask available to user's `compute`.
|
||||
compute = self.compute
|
||||
dates, assets = mask.index, mask.columns
|
||||
out = full(mask.shape, nan, dtype=self.dtype)
|
||||
with self.ctx:
|
||||
# TODO: Consider pre-filtering columns that are all-nan at each
|
||||
@@ -249,5 +289,32 @@ class CustomTermMixin(object):
|
||||
out[idx],
|
||||
*(next(w) for w in windows)
|
||||
)
|
||||
out[~mask.values] = nan
|
||||
out[~mask] = nan
|
||||
return out
|
||||
|
||||
|
||||
class AssetExists(Term):
|
||||
"""
|
||||
Pseudo-filter describing whether or not an asset existed on a given day.
|
||||
This is the default mask for all terms that haven't been passed a mask
|
||||
explicitly.
|
||||
|
||||
This is morally a Filter, in the sense that it produces a boolean value for
|
||||
every asset on every date. We don't subclass Filter, however, because
|
||||
`AssetExists` is computed directly by the FFCEngine.
|
||||
|
||||
See Also
|
||||
--------
|
||||
zipline.assets.AssetFinder.lifetimes
|
||||
"""
|
||||
inputs = ()
|
||||
dtype = bool_
|
||||
window_length = 0
|
||||
mask = None
|
||||
|
||||
def _compute(self, *args, **kwargs):
|
||||
# TODO: Consider moving the bulk of the logic from
|
||||
# SimpleFFCEngine._compute_root_mask here.
|
||||
raise NotImplementedError(
|
||||
"Direct computation of AssetExists is not supported!"
|
||||
)
|
||||
|
||||
@@ -5,14 +5,14 @@ from functools import partial
|
||||
from contextlib import contextmanager
|
||||
from logbook import Logger, StderrHandler
|
||||
from networkx import topological_sort
|
||||
from six import iteritems, itervalues
|
||||
from six import iteritems
|
||||
import subprocess
|
||||
|
||||
from zipline.data.dataset import BoundColumn
|
||||
from zipline.modelling import Filter, Factor, Classifier, Term
|
||||
from zipline.modelling.term import AssetExists
|
||||
from zipline.modelling.graph import TermGraph
|
||||
|
||||
|
||||
logger = Logger('Visualize')
|
||||
|
||||
|
||||
@@ -70,7 +70,7 @@ def roots(g):
|
||||
return set(n for n, d in iteritems(g.in_degree()) if d == 0)
|
||||
|
||||
|
||||
def write_graph(g, filename, formats=('svg',)):
|
||||
def write_graph(g, filename, formats=('svg',), include_asset_exists=False):
|
||||
"""
|
||||
Write the dependency graph of `terms` as a dot graph.
|
||||
|
||||
@@ -95,6 +95,8 @@ def write_graph(g, filename, formats=('svg',)):
|
||||
# Write inputs cluster.
|
||||
with cluster(f, 'Input', **cluster_attrs):
|
||||
for term in in_nodes:
|
||||
if term is AssetExists() and not include_asset_exists:
|
||||
continue
|
||||
add_term_node(f, term)
|
||||
|
||||
# Write intermediate results.
|
||||
@@ -105,6 +107,8 @@ def write_graph(g, filename, formats=('svg',)):
|
||||
|
||||
# Write edges
|
||||
for source, dest in g.edges():
|
||||
if source is AssetExists() and not include_asset_exists:
|
||||
continue
|
||||
add_edge(f, id(source), id(dest))
|
||||
|
||||
outs = []
|
||||
@@ -116,7 +120,7 @@ def write_graph(g, filename, formats=('svg',)):
|
||||
return outs
|
||||
|
||||
|
||||
def show_graph(g):
|
||||
def show_graph(g, include_asset_exists=False):
|
||||
"""
|
||||
Display a TermGraph interactively with IPython
|
||||
"""
|
||||
@@ -124,7 +128,12 @@ def show_graph(g):
|
||||
from IPython.display import SVG
|
||||
except ImportError:
|
||||
raise Exception("IPython is not installed. Can't show term graph.")
|
||||
result = write_graph(g, 'temp', ('svg',))[0]
|
||||
result = write_graph(
|
||||
g,
|
||||
'temp',
|
||||
('svg',),
|
||||
include_asset_exists=include_asset_exists,
|
||||
)[0]
|
||||
return SVG(filename=result)
|
||||
|
||||
|
||||
|
||||
+8
-1
@@ -504,7 +504,14 @@ class BarData(object):
|
||||
if today > self._factor_matrix_expires:
|
||||
self._factor_matrix, self._factor_matrix_expires = \
|
||||
algo.compute_factor_matrix(today)
|
||||
return self._factor_matrix.loc[today]
|
||||
try:
|
||||
return self._factor_matrix.loc[today]
|
||||
except KeyError:
|
||||
# This happens if no assets passed our filters on a given day.
|
||||
return pd.DataFrame(
|
||||
index=[],
|
||||
columns=self._factor_matrix.columns,
|
||||
)
|
||||
|
||||
def __contains__(self, name):
|
||||
if self._contains_override:
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
"""
|
||||
Utilities for working with pandas objects.
|
||||
"""
|
||||
|
||||
|
||||
def explode(df):
|
||||
"""
|
||||
Take a DataFrame and return a triple of
|
||||
|
||||
(df.index, df.columns, df.values)
|
||||
"""
|
||||
return df.index, df.columns, df.values
|
||||
Reference in New Issue
Block a user