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:
Scott Sanderson
2015-09-16 01:47:11 -04:00
parent 691f6d95d7
commit 26fd6fda8b
32 changed files with 1643 additions and 733 deletions
+6 -2
View File
@@ -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
View File
@@ -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
View File
@@ -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)),
)
+47 -5
View File
@@ -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])
+76 -42
View File
@@ -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)
+15 -12
View File
@@ -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)
+192 -19
View File
@@ -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)
+3 -1
View File
@@ -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)
+21 -9
View File
@@ -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
View File
@@ -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
View File
@@ -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,
+19 -3
View File
@@ -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)
+5
View File
@@ -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,
+6 -6
View File
@@ -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),
)]
+282
View File
@@ -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
+196
View File
@@ -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
+3 -186
View File
@@ -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
+82 -58
View File
@@ -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]
+5 -5
View File
@@ -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
+10
View File
@@ -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
+1 -2
View File
@@ -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
View File
@@ -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,
)
)
+3 -3
View File
@@ -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.
"""
+31 -25
View File
@@ -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):
+15
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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!"
)
+14 -5
View File
@@ -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
View File
@@ -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:
+12
View File
@@ -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