wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-11 02:27:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

ENH: Pipeline API

Browse Source 
- Adds `zipline.pipeline.Pipeline`, a new user-facing class for managing
  pipelines of Modeling API expressions.

- Adds `attach_pipeline` and `drain_pipeline` as API methods

- Removes `add_factor` and `add_filter` as API methods.  These have been
  replaced two new methods on `Pipeline`: `add`, and `apply_screen`.

- Adding a `Filter` as a column no longer implicitly truncates rows from
  the Modelling API output.  It simply causes a new column, of dtype
  `bool` to show up in the output. Removal of rows is now handled by the
  new `apply_screen` method of `Pipeline`.

- Refactors the existing Modeling API tests to reflect the new APIs.
This commit is contained in:
Scott Sanderson
2015-09-29 11:31:31 -04:00
parent 0e0dec49e5
commit 8e59d12daf
12 changed files with 782 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/modelling/base.py
+5 -6
View File
@@ -10,7 +10,6 @@ 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
@@ -72,15 +71,15 @@ class BaseFFCTestCase(TestCase):
"""Default shape for methods that build test data."""
return self.__mask.shape
def run_terms(self, terms, initial_workspace, mask=None):
def run_graph(self, graph, initial_workspace, mask=None):
"""
Compute the given terms, seeding the workspace of our FFCEngine with
Compute the given TermGraph, seeding the workspace of our engine with
`initial_workspace`.
Parameters
----------
terms : dict
Mapping from termname -> term object.
graph : zipline.pipeline.graph.TermGraph
Graph to run.
initial_workspace : dict
Initial workspace to forward to SimpleFFCEngine.compute_chunk.
mask : DataFrame, optional
@@ -104,7 +103,7 @@ class BaseFFCTestCase(TestCase):
dates, assets, mask_values = explode(mask)
initial_workspace.setdefault(AssetExists(), mask_values)
return engine.compute_chunk(
TermGraph(terms),
graph,
dates,
assets,
initial_workspace,
tests/modelling/test_engine.py
+107 -38
View File
@@ -6,11 +6,12 @@ from unittest import TestCase
from itertools import product
from numpy import (
array,
full,
nan,
tile,
zeros,
)
from numpy.testing import assert_array_equal
from pandas import (
DataFrame,
date_range,
@@ -45,6 +46,7 @@ from zipline.modelling.factor.technical import (
MaxDrawdown,
SimpleMovingAverage,
)
from zipline.modelling.pipeline import Pipeline
from zipline.utils.memoize import lazyval
from zipline.utils.test_utils import (
make_rotating_asset_info,
@@ -62,6 +64,22 @@ class RollingSumDifference(CustomFactor):
out[:] = (open - close).sum(axis=0)
class AssetID(CustomFactor):
"""
CustomFactor that returns the AssetID of each asset.
Useful for providing a Factor that produces a different value for each
asset.
"""
window_length = 1
# HACK: We currently decide whether to load or compute a Term based on the
# length of its inputs. This means we have to provide a dummy input.
inputs = [USEquityPricing.close]
def compute(self, today, assets, out, close):
out[:] = assets
def assert_multi_index_is_product(testcase, index, *levels):
"""Assert that a MultiIndex contains the product of `*levels`."""
testcase.assertIsInstance(
@@ -102,11 +120,36 @@ class ConstantInputTestCase(TestCase):
loader = self.loader
engine = SimpleFFCEngine(loader, self.dates, self.asset_finder)
p = Pipeline('test')
msg = "start_date must be before end_date .*"
with self.assertRaisesRegexp(ValueError, msg):
engine.factor_matrix({}, self.dates[2], self.dates[1])
engine.run_pipeline(p, self.dates[2], self.dates[1])
with self.assertRaisesRegexp(ValueError, msg):
engine.factor_matrix({}, self.dates[2], self.dates[2])
engine.run_pipeline(p, self.dates[2], self.dates[2])
def test_screen(self):
loader = self.loader
finder = self.asset_finder
assets = array(self.assets)
engine = SimpleFFCEngine(loader, self.dates, self.asset_finder)
num_dates = 5
dates = self.dates[10:10 + num_dates]
factor = AssetID()
for asset in assets:
p = Pipeline('test', columns={'f': factor}, screen=factor <= asset)
result = engine.run_pipeline(p, dates[0], dates[-1])
expected_sids = assets[assets <= asset]
expected_assets = finder.retrieve_all(expected_sids)
expected_result = DataFrame(
index=MultiIndex.from_product([dates, expected_assets]),
data=tile(expected_sids.astype(float), [len(dates)]),
columns=['f'],
)
assert_frame_equal(result, expected_result)
def test_single_factor(self):
loader = self.loader
@@ -117,17 +160,29 @@ class ConstantInputTestCase(TestCase):
dates = self.dates[10:10 + num_dates]
factor = RollingSumDifference()
expected_result = -factor.window_length
result = engine.factor_matrix({'f': factor}, dates[0], dates[-1])
self.assertEqual(set(result.columns), {'f'})
assert_multi_index_is_product(
self, result.index, dates, finder.retrieve_all(assets)
)
# Since every asset will pass the screen, these should be equivalent.
pipelines = [
Pipeline('test', columns={'f': factor}),
Pipeline(
'test',
columns={'f': factor},
screen=factor.eq(expected_result),
),
]
assert_array_equal(
result['f'].unstack().values,
full(result_shape, -factor.window_length),
)
for p in pipelines:
result = engine.run_pipeline(p, dates[0], dates[-1])
self.assertEqual(set(result.columns), {'f'})
assert_multi_index_is_product(
self, result.index, dates, finder.retrieve_all(assets)
)
check_arrays(
result['f'].unstack().values,
full(result_shape, expected_result),
)
def test_multiple_rolling_factors(self):
@@ -145,27 +200,32 @@ class ConstantInputTestCase(TestCase):
inputs=[USEquityPricing.open, USEquityPricing.high],
)
results = engine.factor_matrix(
{'short': short_factor, 'long': long_factor, 'high': high_factor},
dates[0],
dates[-1],
pipeline = Pipeline(
'test',
columns={
'short': short_factor,
'long': long_factor,
'high': high_factor,
}
)
results = engine.run_pipeline(pipeline, dates[0], dates[-1])
self.assertEqual(set(results.columns), {'short', 'high', 'long'})
assert_multi_index_is_product(
self, results.index, dates, finder.retrieve_all(assets)
)
# row-wise sum over an array whose values are all (1 - 2)
assert_array_equal(
check_arrays(
results['short'].unstack().values,
full(shape, -short_factor.window_length),
)
assert_array_equal(
check_arrays(
results['long'].unstack().values,
full(shape, -long_factor.window_length),
)
# row-wise sum over an array whose values are all (1 - 3)
assert_array_equal(
check_arrays(
results['high'].unstack().values,
full(shape, -2 * high_factor.window_length),
)
@@ -183,12 +243,15 @@ class ConstantInputTestCase(TestCase):
open_minus_close = RollingSumDifference(inputs=[open, close])
avg = (high_minus_low + open_minus_close) / 2
results = engine.factor_matrix(
{
'high_low': high_minus_low,
'open_close': open_minus_close,
'avg': avg,
},
results = engine.run_pipeline(
Pipeline(
'test',
columns={
'high_low': high_minus_low,
'open_close': open_minus_close,
'avg': avg,
},
),
dates[0],
dates[-1],
)
@@ -311,8 +374,11 @@ class FrameInputTestCase(TestCase):
)
bounds = product_upper_triangle(range(window_length, len(dates)))
for start, stop in bounds:
results = engine.factor_matrix(
{'low': low_mavg, 'high': high_mavg},
results = engine.run_pipeline(
Pipeline(
'test',
columns={'low': low_mavg, 'high': high_mavg}
),
dates[start],
dates[stop],
)
@@ -424,8 +490,8 @@ class SyntheticBcolzTestCase(TestCase):
window_length=window_length,
)
results = engine.factor_matrix(
{'sma': SMA},
results = engine.run_pipeline(
Pipeline('test', columns={'sma': SMA}),
dates_to_test[0],
dates_to_test[-1],
)
@@ -476,8 +542,8 @@ class SyntheticBcolzTestCase(TestCase):
window_length=window_length,
)
results = engine.factor_matrix(
{'drawdown': drawdown},
results = engine.run_pipeline(
Pipeline('test', columns={'drawdown': drawdown}),
dates_to_test[0],
dates_to_test[-1],
)
@@ -529,13 +595,16 @@ class MultiColumnLoaderTestCase(TestCase):
sumdiff = RollingSumDifference()
result = engine.factor_matrix(
{
'sumdiff': sumdiff,
'open': open_.latest,
'close': close.latest,
'volume': volume.latest,
},
result = engine.run_pipeline(
Pipeline(
'test',
columns={
'sumdiff': sumdiff,
'open': open_.latest,
'close': close.latest,
'volume': volume.latest,
},
),
dates_to_test[0],
dates_to_test[-1]
)
tests/modelling/test_factor.py
+17 -12
View File
@@ -5,6 +5,7 @@ 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.modelling.graph import TermGraph
from zipline.utils.test_utils import check_arrays
from .base import BaseFFCTestCase
@@ -68,8 +69,9 @@ class FactorTestCase(BaseFFCTestCase):
}
def check(terms):
results = self.run_terms(
terms,
graph = TermGraph(terms)
results = self.run_graph(
graph,
initial_workspace={self.f: data},
mask=self.build_mask(ones((5, 5))),
)
@@ -123,8 +125,9 @@ class FactorTestCase(BaseFFCTestCase):
}
def check(terms):
results = self.run_terms(
terms,
graph = TermGraph(terms)
results = self.run_graph(
graph,
initial_workspace={self.f: data},
mask=self.build_mask(ones((5, 5))),
)
@@ -148,12 +151,14 @@ class FactorTestCase(BaseFFCTestCase):
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()),
}
graph = TermGraph(
{
"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.],
@@ -180,8 +185,8 @@ class FactorTestCase(BaseFFCTestCase):
[4., 3., 2., 1., nan]]),
}
results = self.run_terms(
terms,
results = self.run_graph(
graph,
initial_workspace,
mask=self.build_mask(ones((5, 5))),
)
tests/modelling/test_filter.py
+36 -27
View File
@@ -22,6 +22,7 @@ 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.modelling.graph import TermGraph
from zipline.utils.test_utils import check_arrays
from .base import BaseFFCTestCase, with_default_shape
@@ -108,7 +109,7 @@ class FilterTestCase(BaseFFCTestCase):
nan_data[:, 0] = nan
mask = Mask()
initial_workspace = {self.f: data, mask: mask_data}
workspace = {self.f: data, mask: mask_data}
methods = ['top', 'bottom']
counts = 2, 3, 10
@@ -127,7 +128,7 @@ class FilterTestCase(BaseFFCTestCase):
term = getattr(self.f, method)(**kwargs)
terms[termname(method, count, masked)] = term
results = self.run_terms(terms, initial_workspace=initial_workspace)
results = self.run_graph(TermGraph(terms), initial_workspace=workspace)
def expected_result(method, count, masked):
# Ranking with a mask is equivalent to ranking with nans applied on
@@ -155,8 +156,10 @@ class FilterTestCase(BaseFFCTestCase):
def test_bottom(self):
counts = 2, 3, 10
data = self.randn_data(seed=5) # Arbitrary seed choice.
results = self.run_terms(
terms={'bottom_' + str(c): self.f.bottom(c) for c in counts},
results = self.run_graph(
TermGraph(
{'bottom_' + str(c): self.f.bottom(c) for c in counts}
),
initial_workspace={self.f: data},
)
for c in counts:
@@ -179,15 +182,17 @@ class FilterTestCase(BaseFFCTestCase):
filter_names = ['pct_' + str(q) for q in quintiles]
iter_quintiles = zip(filter_names, quintiles)
terms = {
name: self.f.percentile_between(q * 20.0, (q + 1) * 20.0)
for name, q in zip(filter_names, quintiles)
}
graph = TermGraph(
{
name: self.f.percentile_between(q * 20.0, (q + 1) * 20.0)
for name, q in zip(filter_names, quintiles)
}
)
# Test with 5 columns and no NaNs.
eye5 = eye(5, dtype=float64)
results = self.run_terms(
terms,
results = self.run_graph(
graph,
initial_workspace={self.f: eye5},
mask=self.build_mask(ones((5, 5))),
)
@@ -211,8 +216,8 @@ class FilterTestCase(BaseFFCTestCase):
[1, 1, 1, 0, 1, 1],
[1, 1, 1, 1, 0, 1]], dtype=bool)
results = self.run_terms(
terms,
results = self.run_graph(
graph,
initial_workspace={self.f: eye6},
mask=self.build_mask(mask)
)
@@ -231,8 +236,8 @@ class FilterTestCase(BaseFFCTestCase):
# In particular, the NaNs should never pass any filters.
eye6_withnans = eye6.copy()
putmask(eye6_withnans, ~mask, nan)
results = self.run_terms(
terms,
results = self.run_graph(
graph,
initial_workspace={self.f: eye6},
mask=self.build_mask(mask)
)
@@ -258,12 +263,14 @@ class FilterTestCase(BaseFFCTestCase):
quartiles = range(4)
filter_names = ['pct_' + str(q) for q in quartiles]
terms = {
name: self.f.percentile_between(q * 25.0, (q + 1) * 25.0)
for name, q in zip(filter_names, quartiles)
}
results = self.run_terms(
terms,
graph = TermGraph(
{
name: self.f.percentile_between(q * 25.0, (q + 1) * 25.0)
for name, q in zip(filter_names, quartiles)
}
)
results = self.run_graph(
graph,
initial_workspace={self.f: data},
mask=self.build_mask(ones((5, 5))),
)
@@ -287,14 +294,16 @@ class FilterTestCase(BaseFFCTestCase):
without_mask = self.g.percentile_between(80, 100)
with_mask = self.g.percentile_between(80, 100, mask=custom_mask)
terms = {
'custom_mask': custom_mask,
'without': without_mask,
'with': with_mask,
}
graph = TermGraph(
{
'custom_mask': custom_mask,
'without': without_mask,
'with': with_mask,
}
)
results = self.run_terms(
terms,
results = self.run_graph(
graph,
initial_workspace={self.f: f_input, self.g: g_input},
mask=initial_mask,
)
tests/modelling/test_modelling_algo.py
+119 -8
View File
@@ -29,7 +29,8 @@ from testfixtures import TempDirectory
from zipline.algorithm import TradingAlgorithm
from zipline.api import (
add_factor,
attach_pipeline,
drain_pipeline,
get_datetime,
)
from zipline.data.equities import USEquityPricing
@@ -41,9 +42,15 @@ from zipline.data.ffc.loaders.us_equity_pricing import (
SQLiteAdjustmentWriter,
USEquityPricingLoader,
)
from zipline.errors import (
AttachPipelineAfterInitialize,
DrainPipelineDuringInitialize,
NoSuchPipeline,
)
from zipline.finance import trading
from zipline.modelling.factor.technical import VWAP
from zipline.modelling.pipeline import Pipeline
from zipline.utils.test_utils import (
make_simple_asset_info,
str_to_seconds,
@@ -157,26 +164,127 @@ class ClosesOnly(TestCase):
def exists(self, date, asset):
return asset.start_date <= date <= asset.end_date
def test_attach_pipeline_after_initialize(self):
"""
Assert that calling attach_pipeline after initialize raises correctly.
"""
def initialize(context):
pass
def late_attach(context, data):
attach_pipeline(Pipeline('test'))
raise AssertionError("Shouldn't make it past attach_pipeline!")
algo = TradingAlgorithm(
initialize=initialize,
handle_data=late_attach,
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,
)
with self.assertRaises(AttachPipelineAfterInitialize):
algo.run(source=self.closes)
def barf(context, data):
raise AssertionError("Shouldn't make it past before_trading_start")
algo = TradingAlgorithm(
initialize=initialize,
before_trading_start=late_attach,
handle_data=barf,
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,
)
with self.assertRaises(AttachPipelineAfterInitialize):
algo.run(source=self.closes)
def test_drain_pipeline_after_initialize(self):
"""
Assert that calling drain_pipeline after initialize raises correctly.
"""
def initialize(context):
attach_pipeline(Pipeline('test'))
drain_pipeline('test')
raise AssertionError("Shouldn't make it past drain_pipeline()")
def handle_data(context, data):
raise AssertionError("Shouldn't make it past initialize!")
def before_trading_start(context, data):
raise AssertionError("Shouldn't make it past initialize!")
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,
)
with self.assertRaises(DrainPipelineDuringInitialize):
algo.run(source=self.closes)
def test_drain_nonexistent_pipeline(self):
"""
Assert that calling add_pipeline after initialize raises appropriately.
"""
def initialize(context):
attach_pipeline(Pipeline('test'))
def handle_data(context, data):
raise AssertionError("Shouldn't make it past before_trading_start")
def before_trading_start(context, data):
drain_pipeline('not_test')
raise AssertionError("Shouldn't make it past drain_pipeline!")
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,
)
with self.assertRaises(NoSuchPipeline):
algo.run(source=self.closes)
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')
p = Pipeline('test')
p.add(USEquityPricing.close.latest, 'close')
attach_pipeline(p)
def handle_data(context, data):
factors = data.factors
results = drain_pipeline('test')
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']
latest = results.loc[asset, 'close']
self.assertEqual(latest, self.expected_close(date, asset))
else:
self.assertNotIn(asset, factors.index)
self.assertNotIn(asset, results.index)
before_trading_start = handle_data
@@ -355,17 +463,20 @@ class FFCAlgorithmTestCase(TestCase):
)
def initialize(context):
pipeline = Pipeline('test')
context.vwaps = []
for length, key in iteritems(vwap_keys):
context.vwaps.append(VWAP(window_length=length))
add_factor(context.vwaps[-1], name=key)
pipeline.add(context.vwaps[-1], name=key)
attach_pipeline(pipeline)
def handle_data(context, data):
today = get_datetime()
factors = data.factors
results = drain_pipeline('test')
for length, key in iteritems(vwap_keys):
for asset in assets:
computed = factors.loc[asset, key]
computed = results.loc[asset, key]
expected = vwaps[length][asset].loc[today]
# Only having two places of precision here is a bit
tests/modelling/test_pipeline.py
+127
View File
@@ -0,0 +1,127 @@
"""
Tests for zipline.modelling.pipeline.Pipeline
"""
from unittest import TestCase
from zipline.data.equities import USEquityPricing
from zipline.modelling.pipeline import Pipeline
from zipline.modelling.factor import Factor
from zipline.modelling.filter import Filter
class SomeFactor(Factor):
window_length = 5
inputs = [USEquityPricing.close, USEquityPricing.high]
class SomeOtherFactor(Factor):
window_length = 5
inputs = [USEquityPricing.close, USEquityPricing.high]
class SomeFilter(Filter):
window_length = 5
inputs = [USEquityPricing.close, USEquityPricing.high]
class SomeOtherFilter(Filter):
window_length = 5
inputs = [USEquityPricing.close, USEquityPricing.high]
class PipelineTestCase(TestCase):
def test_construction(self):
p0 = Pipeline('arglebargle')
self.assertEqual(p0.name, 'arglebargle')
self.assertEqual(p0.columns, {})
self.assertIs(p0.screen, None)
columns = {'f': SomeFactor()}
p1 = Pipeline('test', columns=columns)
self.assertEqual(p1.columns, columns)
screen = SomeFilter()
p2 = Pipeline('test', screen=screen)
self.assertEqual(p2.columns, {})
self.assertEqual(p2.screen, screen)
p3 = Pipeline('test', columns=columns, screen=screen)
self.assertEqual(p3.columns, columns)
self.assertEqual(p3.screen, screen)
def test_construction_bad_input_types(self):
with self.assertRaises(TypeError):
Pipeline(1)
with self.assertRaises(TypeError):
Pipeline('test', 1)
Pipeline('test', {})
with self.assertRaises(TypeError):
Pipeline('test', {}, 1)
with self.assertRaises(TypeError):
Pipeline('test', {}, SomeFactor())
Pipeline('test', {}, SomeFactor() > 5)
def test_add(self):
p = Pipeline('test')
f = SomeFactor()
p.add(f, 'f')
self.assertEqual(p.columns, {'f': f})
p.add(f > 5, 'g')
self.assertEqual(p.columns, {'f': f, 'g': f > 5})
with self.assertRaises(TypeError):
p.add(f, 1)
def test_overwrite(self):
p = Pipeline('test')
f = SomeFactor()
other_f = SomeOtherFactor()
p.add(f, 'f')
self.assertEqual(p.columns, {'f': f})
with self.assertRaises(KeyError) as e:
p.add(other_f, 'f')
[message] = e.exception.args
self.assertEqual(message, "Column 'f' already exists.")
p.add(other_f, 'f', overwrite=True)
self.assertEqual(p.columns, {'f': other_f})
def test_remove(self):
f = SomeFactor()
p = Pipeline('test', columns={'f': f})
with self.assertRaises(KeyError) as e:
p.remove('not_a_real_name')
self.assertEqual(f, p.remove('f'))
with self.assertRaises(KeyError) as e:
p.remove('f')
self.assertEqual(e.exception.args, ('f',))
def test_set_screen(self):
f, g = SomeFilter(), SomeOtherFilter()
p = Pipeline('test')
self.assertEqual(p.screen, None)
p.set_screen(f)
self.assertEqual(p.screen, f)
with self.assertRaises(ValueError):
p.set_screen(f)
p.set_screen(g, overwrite=True)
self.assertEqual(p.screen, g)
zipline/algorithm.py
+109 -45
View File
@@ -17,6 +17,7 @@ import warnings
import pytz
import pandas as pd
from pandas.tseries.tools import normalize_date
import numpy as np
from datetime import datetime
@@ -33,16 +34,18 @@ from operator import attrgetter
from zipline.errors import (
AddTermPostInit,
AttachPipelineAfterInitialize,
NoSuchPipeline,
OrderDuringInitialize,
OverrideCommissionPostInit,
OverrideSlippagePostInit,
DrainPipelineDuringInitialize,
RegisterAccountControlPostInit,
RegisterTradingControlPostInit,
UnsupportedCommissionModel,
UnsupportedDatetimeFormat,
UnsupportedOrderParameters,
UnsupportedSlippageModel,
UnsupportedDatetimeFormat,
)
from zipline.finance.trading import TradingEnvironment
from zipline.finance.blotter import Blotter
@@ -78,9 +81,11 @@ from zipline.modelling.engine import (
from zipline.sources import DataFrameSource, DataPanelSource
from zipline.utils.api_support import (
api_method,
require_initialized,
require_not_initialized,
ZiplineAPI,
)
from zipline.utils.cache import CachedObject, Expired
import zipline.utils.events
from zipline.utils.events import (
EventManager,
@@ -223,12 +228,13 @@ class TradingAlgorithm(object):
# Pull in the environment's new AssetFinder for quick reference
self.asset_finder = self.trading_environment.asset_finder
self.init_engine(kwargs.pop('ffc_loader', None))
# Maps from name to Term
self._filters = {}
self._factors = {}
self._classifiers = {}
# Initialize Modeling API data.
self.init_engine(kwargs.pop('ffc_loader', None))
self._pipelines = []
# Create an always-expired cache so that we compute the first time data
# is requested.
self._pipeline_cache = CachedObject(None, pd.Timestamp(0, tz='UTC'))
self.blotter = kwargs.pop('blotter', None)
if not self.blotter:
@@ -1326,41 +1332,96 @@ class TradingAlgorithm(object):
"""
self.register_trading_control(LongOnly())
###########
# FFC API #
###########
##############
# Modeling API
##############
@api_method
@require_not_initialized(AddTermPostInit())
def add_factor(self, factor, name):
if name in self._factors:
raise ValueError("Name %r is already a factor!" % name)
self._factors[name] = factor
@api_method
@require_not_initialized(AddTermPostInit())
def add_filter(self, filter):
name = "anon_filter_%d" % len(self._filters)
self._filters[name] = filter
# Note: add_classifier is not yet implemented since you can't do anything
# useful with classifiers yet.
def _all_terms(self):
# Merge all three dicts.
return dict(
chain.from_iterable(
iteritems(terms)
for terms in (self._filters, self._factors, self._classifiers)
)
)
def compute_factor_matrix(self, start_date):
@require_not_initialized(AttachPipelineAfterInitialize())
def attach_pipeline(self, pipeline):
"""
Compute a factor matrix containing at least the data necessary to
provide values for `start_date`.
Register a pipeline to be computed at the start of each day.
"""
if self._pipelines:
raise NotImplementedError("Multiple pipelines are not supported.")
self._pipelines.append(pipeline)
Loads a factor matrix with data extending from `start_date` until a
year from `start_date`, or until the end of the simulation.
@api_method
@require_initialized(DrainPipelineDuringInitialize())
def drain_pipeline(self, name=None):
"""
Get the results of pipeline with name `name`.
Parameters
----------
name : str or None
Name of the pipeline for which results are requested.
Returns
-------
results : pd.DataFrame
DataFrame containing the results of the requested pipeline for
the current simulation date.
Raises
------
NoSuchPipeline
Raised when no pipeline with the name `name` has been registered.
See Also
--------
:meth:`zipline.modelling.FFCEngine.run_pipeline`
"""
# NOTE: We don't currently support multiple pipelines, but we plan to
# in the future.
for p in self._pipelines:
if p.name == name:
break
# This is a for-else block. Yes, that's a thing in Python.
else:
raise NoSuchPipeline(
name=name,
valid=[p.name for p in self._pipelines],
)
return self._pipeline_results(p)
def _pipeline_results(self, pipeline):
"""
Internal implementation of `drain_pipeline`.
"""
today = normalize_date(self.get_datetime())
try:
data = self._pipeline_cache.unwrap(today)
except Expired:
data, valid_until = self._run_pipeline(pipeline, today)
self._pipeline_cache = CachedObject(data, valid_until)
# Now that we have a cached result, try to return the data for today.
try:
return data.loc[today]
except KeyError:
# This happens if no assets passed the pipeline screen on a given
# day.
return pd.DataFrame(index=[], columns=data.columns)
def _run_pipeline(self, pipeline, start_date):
"""
Compute `pipeline`, providing values for at least `start_date`.
Produces a DataFrame containing data for days between `start_date` and
`end_date`, where `end_date` is defined by:
`end_date = min(start_date + 252 trading days, simulation_end)`
252 is a mostly-arbitrary number based on napkin math. The window
length will likely become dynamic and/or configurable in the future.
Returns
-------
(data, valid_until) : tuple (pd.DataFrame, pd.Timestamp)
See Also
--------
FFCEngine.run_pipeline
"""
days = self.trading_environment.trading_days
@@ -1369,16 +1430,19 @@ class TradingAlgorithm(object):
# ...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.
# choice that seemed reasonable based on napkin math. In the future,
# this number will likely become dynamic and/or customizable, so don't
# rely on it being 252.
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,
end_date,
), end_date
return \
self.engine.run_pipeline(pipeline, start_date, end_date), end_date
##################
# End Modeling API
##################
def current_universe(self):
return self._current_universe
zipline/errors.py
+24 -6
View File
@@ -377,15 +377,33 @@ class UnknownRankMethod(ZiplineError):
)
class AddTermPostInit(ZiplineError):
class AttachPipelineAfterInitialize(ZiplineError):
"""
Raised when a user tries to call add_{filter,factor,classifier}
outside of initialize.
Raised when a user tries to call add_pipeline outside of initialize.
"""
msg = (
"Attempted to add a new filter, factor, or classifier "
"outside of initialize.\n"
"New FFC terms may only be added during initialize."
"Attempted to attach a pipeline after initialize()."
"attach_pipeline() can only be called during initialize."
)
class DrainPipelineDuringInitialize(ZiplineError):
"""
Raised when a user tries to call `drain_pipeline` during initialize.
"""
msg = (
"Attempted to call drain_pipeline() during initialize. "
"drain_pipeline() can only be called once initialize has completed."
)
class NoSuchPipeline(ZiplineError, KeyError):
"""
Raised when a user tries to access a non-existent pipeline by name.
"""
msg = (
"No pipeline named '{name}' exists. Valid pipeline names are {valid}. "
"Did you forget to call attach_pipeline()?"
)
zipline/modelling/engine.py
+59 -135
View File
@@ -5,21 +5,15 @@ from abc import (
ABCMeta,
abstractmethod,
)
from operator import and_
from uuid import uuid4
from six import (
iteritems,
itervalues,
with_metaclass,
)
from six.moves import (
reduce,
zip_longest,
)
from six.moves import zip_longest
from numpy import array
from numpy import (
add,
empty_like,
)
from pandas import (
DataFrame,
date_range,
@@ -28,32 +22,25 @@ from pandas import (
from zipline.lib.adjusted_array import ensure_ndarray
from zipline.errors import NoFurtherDataError
from zipline.utils.numpy_utils import repeat_first_axis, repeat_last_axis
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)):
@abstractmethod
def factor_matrix(self, terms, start_date, end_date):
def run_pipeline(self, pipeline, start_date, end_date):
"""
Compute values for `terms` between `start_date` and `end_date`.
Compute values for `pipeline` between `start_date` and `end_date`.
Returns a DataFrame with a MultiIndex of (date, asset) pairs on the
index. On each date, we return a row for each asset that passed all
instances of `Filter` in `terms, and the columns of the returned frame
will be the keys in `terms` whose values are instances of `Factor`.
Returns a DataFrame with a MultiIndex of (date, asset) pairs
Parameters
----------
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.
pipeline : zipline.modelling.pipeline.Pipeline
The pipeline to run.
start_date : pd.Timestamp
Start date of the computed matrix.
end_date : pd.Timestamp
@@ -61,23 +48,31 @@ class FFCEngine(with_metaclass(ABCMeta)):
Returns
-------
matrix : pd.DataFrame
A matrix of computed results.
result : pd.DataFrame
A frame of computed results.
The columns `result` correspond wil be the computed results of
`pipeline.columns`, which should be a dictionary mapping strings to
instances of `zipline.modelling.term.Term`.
For each date between `start_date` and `end_date`, `result` will
contain a row for each asset that passed `pipeline.screen`. A
screen of None indicates that a row should be returned for each
asset that existed each day.
"""
raise NotImplementedError("factor_matrix")
raise NotImplementedError("run_pipeline")
class NoOpFFCEngine(FFCEngine):
"""
FFCEngine that doesn't do anything.
An FFCEngine that doesn't do anything.
"""
def factor_matrix(self, terms, start_date, end_date):
def run_pipeline(self, pipeline, start_date, end_date):
return DataFrame(
index=MultiIndex.from_product(
[date_range(start=start_date, end=end_date, freq='D'), ()],
),
columns=sorted(terms.keys())
columns=sorted(pipeline.columns.keys()),
)
@@ -110,15 +105,14 @@ class SimpleFFCEngine(object):
self._finder = asset_finder
self._root_mask_term = AssetExists()
def factor_matrix(self, terms, start_date, end_date):
def run_pipeline(self, pipeline, start_date, end_date):
"""
Compute a factor matrix.
Compute a pipeline.
Parameters
----------
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.
pipeline : zipline.modelling.pipeline.Pipeline
The pipeline to run.
start_date : pd.Timestamp
Start date of the computed matrix.
end_date : pd.Timestamp
@@ -155,7 +149,7 @@ class SimpleFFCEngine(object):
See Also
--------
FFCEngine.factor_matrix
FFCEngine.run_pipeline
"""
if end_date <= start_date:
raise ValueError(
@@ -163,36 +157,23 @@ class SimpleFFCEngine(object):
"start_date=%s, end_date=%s" % (start_date, end_date)
)
graph = TermGraph(terms)
screen_name = uuid4().hex
graph = pipeline.to_graph(screen_name, self._root_mask_term)
extra_rows = graph.extra_rows[self._root_mask_term]
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(
outputs = self.compute_chunk(
graph,
dates,
assets,
initial_workspace={self._root_mask_term: root_mask_values},
)
# Collect the results that we'll actually show to the user.
filters, factors = {}, {}
for name, term in iteritems(terms):
if isinstance(term, Filter):
filters[name] = raw_outputs[name]
elif isinstance(term, Factor):
factors[name] = raw_outputs[name]
elif isinstance(term, Classifier):
continue
else:
raise ValueError("Unknown term type: %s" % term)
# 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:]
screen_values = outputs.pop(screen_name)
return self._format_factor_matrix(out_dates, assets, filters, factors)
return self._to_narrow(outputs, screen_values, out_dates, assets)
def _compute_root_mask(self, start_date, end_date, extra_rows):
"""
@@ -360,98 +341,41 @@ class SimpleFFCEngine(object):
out[name] = workspace[term][graph_extra_rows[term]:]
return out
def _format_factor_matrix(self, dates, assets, filters, factors):
def _to_narrow(self, data, mask, dates, assets):
"""
Convert raw computed filters/factors into a DataFrame for public APIs.
Convert raw computed pipeline results into a DataFrame for public APIs.
Parameters
----------
dates : np.array[datetime64]
Row index for arrays in `filters` and `factors.`
assets : np.array[int64]
Column index for arrays in `filters` and `factors.`
filters : dict
Dict mapping filter names -> computed filters.
factors : dict
Dict mapping factor names -> computed factors.
data : dict[str -> ndarray[ndim=2]]
Dict mapping column names to computed results.
mask : ndarray[bool, ndim=2]
Mask array of values to keep.
dates : ndarray[datetime64, ndim=1]
Row index for arrays `data` and `mask`
assets : ndarray[int64, ndim=2]
Column index for arrays `data` and `mask`
Returns
-------
factor_matrix : pd.DataFrame
The indices of `factor_matrix` are as follows:
results : pd.DataFrame
The indices of `results` are as follows:
index : two-tiered MultiIndex of (date, asset).
For each date, we return a row for each asset that passed all
filters on that date.
columns : keys from `factor_data`
Contains an entry for each (date, asset) pair corresponding to
a `True` value in `mask`.
columns : Index of str
One column per entry in `data`.
Each date/asset/factor triple contains the computed value of the given
factor on the given date for the given asset.
If mask[date, asset] is True, then result.loc[(date, asset), colname]
will contain the value of data[colname][date, asset].
"""
# FUTURE OPTIMIZATION: Cythonize all of this.
# Boolean mask of values that passed all filters.
unioned = reduce(and_, itervalues(filters))
# Parallel arrays of (x,y) coords for (date, asset) pairs that passed
# all filters. Each entry here will correspond to a row in our output
# frame.
nonzero_xs, nonzero_ys = unioned.nonzero()
# Raw arrays storing (date, asset) pairs.
# These will form the index of our output frame.
raw_dates_index = empty_like(nonzero_xs, dtype='datetime64[ns]')
raw_assets_index = empty_like(nonzero_xs, dtype=int)
# Mapping from column_name -> array.
# This will be the `data` arg to our output frame.
columns = {
name: empty_like(nonzero_xs, dtype=factor.dtype)
for name, factor in iteritems(factors)
}
# We're going to iterate over `iteritems(columns)` a whole bunch of
# times down below. It's faster to construct iterate over a tuple of
# pairs.
columns_iter = tuple(iteritems(columns))
# This is tricky.
# unioned.sum(axis=1) gives us an array of the same size as `dates`
# containing, for each date, the number of assets that passed our
# filters on that date.
# Running this through add.accumulate gives us an array containing, for
# each date, the running total of the number of assets that passed our
# filters on or before that date.
# This means that (bounds[i - 1], bounds[i]) gives us the indices of
# the first and last rows in our output frame for each date in `dates`.
bounds = add.accumulate(unioned.sum(axis=1))
day_start = 0
for day_idx, day_end in enumerate(bounds):
day_bounds = slice(day_start, day_end)
column_indices = nonzero_ys[day_bounds]
raw_dates_index[day_bounds] = dates[day_idx]
raw_assets_index[day_bounds] = assets[column_indices]
for name, colarray in columns_iter:
colarray[day_bounds] = factors[name][day_idx, column_indices]
# Upper bound of current row becomes lower bound for next row.
day_start = day_end
resolved_assets = array(self._finder.retrieve_all(assets))
dates_kept = repeat_last_axis(dates.values, len(assets))[mask]
assets_kept = repeat_first_axis(resolved_assets, len(dates))[mask]
return DataFrame(
data=columns,
index=MultiIndex.from_arrays(
[
raw_dates_index,
# FUTURE OPTIMIZATION:
# Avoid duplicate lookups by grouping and only looking up
# each unique sid once.
self._finder.retrieve_all(raw_assets_index),
],
)
data={name: arr[mask] for name, arr in iteritems(data)},
index=MultiIndex.from_arrays([dates_kept, assets_kept]),
).tz_localize('UTC', level=0)
def _validate_compute_chunk_params(self, dates, assets, initial_workspace):
zipline/modelling/pipeline.py
+157
View File
@@ -0,0 +1,157 @@
from zipline.utils.preprocess import expect_types, optional
from zipline.modelling.term import Term
from zipline.modelling.filter import Filter
from zipline.modelling.graph import TermGraph
class Pipeline(object):
"""
Parameters
----------
name : str, optional
Name for this pipeline.
columns : dict, optional
Initial columns.
screen : zipline.modelling.term.Filter, optional
Initial screen.
Methods
-------
add
remove
apply_screen
Attributes
----------
columns
screen
"""
__slots__ = ('_name', '_columns', '_screen', '__weakref__')
@expect_types(
name=str,
columns=optional(dict),
screen=optional(Filter),
)
def __init__(self, name, columns=None, screen=None):
self._name = name
if columns is None:
columns = {}
self._columns = columns
self._screen = screen
@property
def name(self):
"""
The name of this pipeline.
"""
return self._name
@property
def columns(self):
"""
The columns currently applied to this pipeline.
"""
return self._columns
@property
def screen(self):
"""
The screen applied to the rows of this pipeline.
"""
return self._screen
@expect_types(term=Term, name=str)
def add(self, term, name, overwrite=False):
"""
Add a column.
The results of computing `term` will show up as a column in the
DataFrame produced by running this pipeline.
Parameters
----------
column : zipline.modelling.Term
A Filter, Factor, or Classifier to add to the pipeline.
name : str
Name of the column to add.
overwrite : bool
Whether to overwrite the existing entry if we already have a column
named `name`.
"""
columns = self.columns
if name in columns:
if overwrite:
self.remove(name)
else:
raise KeyError("Column '{}' already exists.".format(name))
self._columns[name] = term
@expect_types(name=str)
def remove(self, name):
"""
Remove a column.
Parameters
----------
name : str
The name of the column to remove.
Raises
------
KeyError
If `name` is not in self.columns.
Returns
-------
removed : zipline.modelling.term.Term
The removed term.
"""
return self.columns.pop(name)
@expect_types(screen=Filter)
def set_screen(self, screen, overwrite=False):
"""
Apply a screen to this Pipeline.
If no screen has yet been applied to the pipeline, this method sets
`screen` as the current screen.
Parameter
---------
filter : zipline.modelling.filter.Filter
The screen to apply.
overwrite : bool
Whether to overwrite any existing screen. If overwrite is False
and self.screen is not None, we raise an error.
"""
if self._screen is not None and not overwrite:
raise ValueError(
"set_screen() called with overwrite=False and screen already "
"set.\n"
"If you want to apply multiple filters as a screen use "
"set_screen(filter1 & filter2 & ...).\n"
"If you want to replace the previous screen with a new one, "
"use set_screen(new_filter, overwrite=True)."
)
self._screen = screen
def to_graph(self, screen_name, default_screen):
"""
Compile into a TermGraph.
Parameters
----------
screen_name : str
Name to supply for self.screen.
default_screen : zipline.modelling.term.Term
Term to use as a screen if self.screen is None.
"""
columns = self.columns.copy()
screen = self.screen
if screen is None:
screen = default_screen
columns[screen_name] = screen
return TermGraph(columns)
zipline/protocol.py
-19
View File
@@ -17,7 +17,6 @@ from copy import copy
from six import iteritems, iterkeys
import pandas as pd
from pandas.tseries.tools import normalize_date
import numpy as np
from . utils.protocol_utils import Enum
@@ -494,24 +493,6 @@ class BarData(object):
def __init__(self, data=None):
self._data = data or {}
self._contains_override = None
self._factor_matrix = None
self._factor_matrix_expires = pd.Timestamp(0, tz='UTC')
@property
def factors(self):
algo = get_algo_instance()
today = normalize_date(algo.get_datetime())
if today > self._factor_matrix_expires:
self._factor_matrix, self._factor_matrix_expires = \
algo.compute_factor_matrix(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:
zipline/utils/api_support.py
+22
View File
@@ -76,3 +76,25 @@ def require_not_initialized(exception):
return method(self, *args, **kwargs)
return wrapped_method
return decorator
def require_initialized(exception):
"""
Decorator for API methods that should only be called after
TradingAlgorithm.initialize. `exception` will be raised if the method is
called before initialize has completed.
Usage
-----
@require_initialized(SomeException("Don't do that!"))
def method(self):
# Do stuff that should only be allowed after initialize.
"""
def decorator(method):
@wraps(method)
def wrapped_method(self, *args, **kwargs):
if not self.initialized:
raise exception
return method(self, *args, **kwargs)
return wrapped_method
return decorator