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

Merge pull request #1811 from quantopian/run-chunked-pipeline

Browse Source 
Run chunked pipeline
This commit is contained in:
Ana Ruelas
2017-06-02 17:29:15 -04:00
committed by GitHub
parent 1ab0cd7113 ad3c72a8fb
commit e6fe3a0e7d
10 changed files with 430 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/test_engine.py
+34
View File
@@ -51,6 +51,7 @@ from zipline.pipeline.factors import (
ExponentialWeightedMovingAverage,
ExponentialWeightedMovingStdDev,
MaxDrawdown,
Returns,
SimpleMovingAverage,
)
from zipline.pipeline.loaders.equity_pricing_loader import (
@@ -77,6 +78,7 @@ from zipline.testing import (
)
from zipline.testing.fixtures import (
WithAdjustmentReader,
WithEquityPricingPipelineEngine,
WithSeededRandomPipelineEngine,
WithTradingEnvironment,
ZiplineTestCase,
@@ -1497,3 +1499,35 @@ class PopulateInitialWorkspaceTestCase(WithConstantInputs, ZiplineTestCase):
precomputed_term_value,
),
)
class ChunkedPipelineTestCase(WithEquityPricingPipelineEngine,
ZiplineTestCase):
PIPELINE_START_DATE = Timestamp('2006-01-05', tz='UTC')
END_DATE = Timestamp('2006-12-29', tz='UTC')
def test_run_chunked_pipeline(self):
"""
Test that running a pipeline in chunks produces the same result as if
it were run all at once
"""
pipe = Pipeline(
columns={
'close': USEquityPricing.close.latest,
'returns': Returns(window_length=2),
'categorical': USEquityPricing.close.latest.quantiles(5)
},
)
pipeline_result = self.pipeline_engine.run_pipeline(
pipe,
start_date=self.PIPELINE_START_DATE,
end_date=self.END_DATE,
)
chunked_result = self.pipeline_engine.run_chunked_pipeline(
pipeline=pipe,
start_date=self.PIPELINE_START_DATE,
end_date=self.END_DATE,
chunksize=22
)
self.assertTrue(chunked_result.equals(pipeline_result))
tests/utils/test_date_utils.py
+86
View File
@@ -0,0 +1,86 @@
from pandas import Timestamp
from nose_parameterized import parameterized
from zipline.testing import ZiplineTestCase
from zipline.utils.calendars import get_calendar
from zipline.utils.date_utils import compute_date_range_chunks
def T(s):
"""
Helpful function to improve readibility.
"""
return Timestamp(s, tz='UTC')
class TestDateUtils(ZiplineTestCase):
@classmethod
def init_class_fixtures(cls):
super(TestDateUtils, cls).init_class_fixtures()
cls.calendar = get_calendar('NYSE')
@parameterized.expand([
(None, [(T('2017-01-03'), T('2017-01-31'))]),
(10, [
(T('2017-01-03'), T('2017-01-17')),
(T('2017-01-18'), T('2017-01-31'))
]),
(15, [
(T('2017-01-03'), T('2017-01-24')),
(T('2017-01-25'), T('2017-01-31'))
]),
])
def test_compute_date_range_chunks(self, chunksize, expected):
# This date range results in 20 business days
start_date = T('2017-01-03')
end_date = T('2017-01-31')
date_ranges = compute_date_range_chunks(
self.calendar.all_sessions,
start_date,
end_date,
chunksize
)
self.assertListEqual(list(date_ranges), expected)
def test_compute_date_range_chunks_invalid_input(self):
# Start date not found in calendar
with self.assertRaises(KeyError) as cm:
compute_date_range_chunks(
self.calendar.all_sessions,
T('2017-05-07'), # Sunday
T('2017-06-01'),
None
)
self.assertEqual(
str(cm.exception),
"'Start date 2017-05-07 is not found in calendar.'"
)
# End date not found in calendar
with self.assertRaises(KeyError) as cm:
compute_date_range_chunks(
self.calendar.all_sessions,
T('2017-05-01'),
T('2017-05-27'), # Saturday
None
)
self.assertEqual(
str(cm.exception),
"'End date 2017-05-27 is not found in calendar.'"
)
# End date before start date
with self.assertRaises(ValueError) as cm:
compute_date_range_chunks(
self.calendar.all_sessions,
T('2017-06-01'),
T('2017-05-01'),
None
)
self.assertEqual(
str(cm.exception),
"End date 2017-05-01 cannot precede start date 2017-06-01."
)
tests/utils/test_pandas_utils.py
+107 -1
View File
@@ -4,7 +4,11 @@ Tests for zipline/utils/pandas_utils.py
import pandas as pd
from zipline.testing import parameter_space, ZiplineTestCase
from zipline.utils.pandas_utils import nearest_unequal_elements
from zipline.testing.predicates import assert_equal
from zipline.utils.pandas_utils import (
categorical_df_concat,
nearest_unequal_elements
)
class TestNearestUnequalElements(ZiplineTestCase):
@@ -80,3 +84,105 @@ class TestNearestUnequalElements(ZiplineTestCase):
str(e.exception),
'dts must be sorted in increasing order',
)
class TestCatDFConcat(ZiplineTestCase):
def test_categorical_df_concat(self):
inp = [
pd.DataFrame(
{
'A': pd.Series(['a', 'b', 'c'], dtype='category'),
'B': pd.Series([100, 102, 103], dtype='int64'),
'C': pd.Series(['x', 'x', 'x'], dtype='category'),
}
),
pd.DataFrame(
{
'A': pd.Series(['c', 'b', 'd'], dtype='category'),
'B': pd.Series([103, 102, 104], dtype='int64'),
'C': pd.Series(['y', 'y', 'y'], dtype='category'),
}
),
pd.DataFrame(
{
'A': pd.Series(['a', 'b', 'd'], dtype='category'),
'B': pd.Series([101, 102, 104], dtype='int64'),
'C': pd.Series(['z', 'z', 'z'], dtype='category'),
}
),
]
result = categorical_df_concat(inp)
expected = pd.DataFrame(
{
'A': pd.Series(
['a', 'b', 'c', 'c', 'b', 'd', 'a', 'b', 'd'],
dtype='category'
),
'B': pd.Series(
[100, 102, 103, 103, 102, 104, 101, 102, 104],
dtype='int64'
),
'C': pd.Series(
['x', 'x', 'x', 'y', 'y', 'y', 'z', 'z', 'z'],
dtype='category'
),
},
)
expected.index = pd.Int64Index([0, 1, 2, 0, 1, 2, 0, 1, 2])
assert_equal(expected, result)
assert_equal(
expected['A'].cat.categories,
result['A'].cat.categories
)
assert_equal(
expected['C'].cat.categories,
result['C'].cat.categories
)
def test_categorical_df_concat_value_error(self):
mismatched_dtypes = [
pd.DataFrame(
{
'A': pd.Series(['a', 'b', 'c'], dtype='category'),
'B': pd.Series([100, 102, 103], dtype='int64'),
}
),
pd.DataFrame(
{
'A': pd.Series(['c', 'b', 'd'], dtype='category'),
'B': pd.Series([103, 102, 104], dtype='float64'),
}
),
]
mismatched_column_names = [
pd.DataFrame(
{
'A': pd.Series(['a', 'b', 'c'], dtype='category'),
'B': pd.Series([100, 102, 103], dtype='int64'),
}
),
pd.DataFrame(
{
'A': pd.Series(['c', 'b', 'd'], dtype='category'),
'X': pd.Series([103, 102, 104], dtype='int64'),
}
),
]
with self.assertRaises(ValueError) as cm:
categorical_df_concat(mismatched_dtypes)
self.assertEqual(
str(cm.exception),
"Input DataFrames must have the same columns/dtypes."
)
with self.assertRaises(ValueError) as cm:
categorical_df_concat(mismatched_column_names)
self.assertEqual(
str(cm.exception),
"Input DataFrames must have the same columns/dtypes."
)
tests/utils/test_sharedoc.py
+21
View File
@@ -0,0 +1,21 @@
from zipline.testing import ZiplineTestCase
from zipline.utils.sharedoc import copydoc
class TestSharedoc(ZiplineTestCase):
def test_copydoc(self):
def original_docstring_function():
"""
My docstring brings the boys to the yard.
"""
pass
@copydoc(original_docstring_function)
def copied_docstring_function():
pass
self.assertEqual(
original_docstring_function.__doc__,
copied_docstring_function.__doc__
)
zipline/pipeline/engine.py
+64 -3
View File
@@ -27,6 +27,10 @@ from zipline.utils.pandas_utils import explode
from .term import AssetExists, InputDates, LoadableTerm
from zipline.utils.date_utils import compute_date_range_chunks
from zipline.utils.pandas_utils import categorical_df_concat
from zipline.utils.sharedoc import copydoc
class PipelineEngine(with_metaclass(ABCMeta)):
@@ -62,6 +66,45 @@ class PipelineEngine(with_metaclass(ABCMeta)):
"""
raise NotImplementedError("run_pipeline")
@abstractmethod
def run_chunked_pipeline(self, pipeline, start_date, end_date, chunksize):
"""
Compute values for `pipeline` in number of days equal to `chunksize`
and return stitched up result. Computing in chunks is useful for
pipelines computed over a long period of time.
Parameters
----------
pipeline : Pipeline
The pipeline to run.
start_date : pd.Timestamp
The start date to run the pipeline for.
end_date : pd.Timestamp
The end date to run the pipeline for.
chunksize : int or None
The number of days to execute at a time. If None, then
results will be calculated for entire date range at once.
Returns
-------
result : pd.DataFrame
A frame of computed results.
The columns `result` correspond to the entries of
`pipeline.columns`, which should be a dictionary mapping strings to
instances of `zipline.pipeline.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.
See Also
--------
:meth:`PipelineEngine.run_pipeline`
"""
raise NotImplementedError("run_chunked_pipeline")
class NoEngineRegistered(Exception):
"""
@@ -80,6 +123,12 @@ class ExplodingPipelineEngine(PipelineEngine):
"resources were registered."
)
def run_chunked_pipeline(self, pipeline, start_date, end_date, chunksize):
raise NoEngineRegistered(
"Attempted to run a chunked pipeline but no pipeline "
"resources were registered."
)
def default_populate_initial_workspace(initial_workspace,
root_mask_term,
@@ -114,7 +163,7 @@ def default_populate_initial_workspace(initial_workspace,
return initial_workspace
class SimplePipelineEngine(object):
class SimplePipelineEngine(PipelineEngine):
"""
PipelineEngine class that computes each term independently.
@@ -146,7 +195,6 @@ class SimplePipelineEngine(object):
'_root_mask_term',
'_root_mask_dates_term',
'_populate_initial_workspace',
'__weakref__',
)
def __init__(self,
@@ -210,7 +258,8 @@ class SimplePipelineEngine(object):
See Also
--------
PipelineEngine.run_pipeline
:meth:`PipelineEngine.run_pipeline`
:meth:`PipelineEngine.run_chunked_pipeline`
"""
if end_date < start_date:
raise ValueError(
@@ -256,6 +305,18 @@ class SimplePipelineEngine(object):
assets,
)
@copydoc(PipelineEngine.run_chunked_pipeline)
def run_chunked_pipeline(self, pipeline, start_date, end_date, chunksize):
ranges = compute_date_range_chunks(
self._calendar,
start_date,
end_date,
chunksize,
)
chunks = [self.run_pipeline(pipeline, s, e) for s, e in ranges]
return categorical_df_concat(chunks, inplace=True)
def _compute_root_mask(self, start_date, end_date, extra_rows):
"""
Compute a lifetimes matrix from our AssetFinder, then drop columns that
zipline/testing/fixtures.py
+7 -5
View File
@@ -1,4 +1,3 @@
from itertools import repeat
import os
import sqlite3
from unittest import TestCase
@@ -1333,12 +1332,15 @@ class WithEquityPricingPipelineEngine(WithAdjustmentReader,
cls.bcolz_equity_daily_bar_reader,
SQLiteAdjustmentReader(cls.adjustments_db_path),
)
dispatcher = dict(
zip(USEquityPricing.columns, repeat(loader))
).__getitem__
def get_loader(column):
if column in USEquityPricing.columns:
return loader
else:
raise AssertionError("No loader registered for %s" % column)
cls.pipeline_engine = SimplePipelineEngine(
get_loader=dispatcher,
get_loader=get_loader,
calendar=cls.nyse_sessions,
asset_finder=cls.asset_finder,
)
zipline/utils/calendars/__init__.py
+6 -6
View File
@@ -15,20 +15,20 @@
from .trading_calendar import TradingCalendar
from .calendar_utils import (
get_calendar,
register_calendar_alias,
register_calendar,
register_calendar_type,
clear_calendars,
deregister_calendar,
clear_calendars
get_calendar,
register_calendar,
register_calendar_alias,
register_calendar_type,
)
__all__ = [
'TradingCalendar',
'clear_calendars',
'deregister_calendar',
'get_calendar',
'register_calendar',
'register_calendar_alias',
'register_calendar_type',
'TradingCalendar',
]
zipline/utils/date_utils.py
+42
View File
@@ -0,0 +1,42 @@
from toolz import partition_all
def compute_date_range_chunks(sessions, start_date, end_date, chunksize):
"""Compute the start and end dates to run a pipeline for.
Parameters
----------
sessions : DatetimeIndex
The available dates.
start_date : pd.Timestamp
The first date in the pipeline.
end_date : pd.Timestamp
The last date in the pipeline.
chunksize : int or None
The size of the chunks to run. Setting this to None returns one chunk.
Returns
-------
ranges : iterable[(np.datetime64, np.datetime64)]
A sequence of start and end dates to run the pipeline for.
"""
if start_date not in sessions:
raise KeyError("Start date %s is not found in calendar." %
(start_date.strftime("%Y-%m-%d"),))
if end_date not in sessions:
raise KeyError("End date %s is not found in calendar." %
(end_date.strftime("%Y-%m-%d"),))
if end_date < start_date:
raise ValueError("End date %s cannot precede start date %s." %
(end_date.strftime("%Y-%m-%d"),
start_date.strftime("%Y-%m-%d")))
if chunksize is None:
return [(start_date, end_date)]
start_ix, end_ix = sessions.slice_locs(start_date, end_date)
return (
(r[0], r[-1]) for r in partition_all(
chunksize, sessions[start_ix:end_ix]
)
)
zipline/utils/pandas_utils.py
+44
View File
@@ -2,6 +2,7 @@
Utilities for working with pandas objects.
"""
from contextlib import contextmanager
from copy import deepcopy
from itertools import product
import operator as op
import warnings
@@ -222,3 +223,46 @@ def clear_dataframe_indexer_caches(df):
delattr(df, attr)
except AttributeError:
pass
def categorical_df_concat(df_list, inplace=False):
"""
Prepare list of pandas DataFrames to be used as input to pd.concat.
Ensure any columns of type 'category' have the same categories across each
dataframe.
Parameters
----------
df_list : list
List of dataframes with same columns.
inplace : bool
True if input list can be modified. Default is False.
Returns
-------
concatenated : df
Dataframe of concatenated list.
"""
if not inplace:
df_list = deepcopy(df_list)
# Assert each dataframe has the same columns/dtypes
df = df_list[0]
if not all([(df.dtypes.equals(df_i.dtypes)) for df_i in df_list[1:]]):
raise ValueError("Input DataFrames must have the same columns/dtypes.")
categorical_columns = df.columns[df.dtypes == 'category']
for col in categorical_columns:
new_categories = sorted(
set().union(
*(frame[col].cat.categories for frame in df_list)
)
)
with ignore_pandas_nan_categorical_warning():
for df in df_list:
df[col].cat.set_categories(new_categories, inplace=True)
return pd.concat(df_list)
zipline/utils/sharedoc.py
+19
View File
@@ -5,6 +5,7 @@ across different functions.
import re
from six import iteritems
from textwrap import dedent
from toolz import curry
PIPELINE_DOWNSAMPLING_FREQUENCY_DOC = dedent(
"""\
@@ -98,3 +99,21 @@ def templated_docstring(**docs):
f.__doc__ = format_docstring(f.__name__, f.__doc__, docs)
return f
return decorator
@curry
def copydoc(from_, to):
"""Copies the docstring from one function to another.
Parameters
----------
from_ : any
The object to copy the docstring from.
to : any
The object to copy the docstring to.
Returns
-------
to : any
``to`` with the docstring from ``from_``
"""
to.__doc__ = from_.__doc__
return to