diff --git a/tests/pipeline/test_run_chunked_pipeline.py b/tests/pipeline/test_run_chunked_pipeline.py new file mode 100644 index 00000000..533ccf06 --- /dev/null +++ b/tests/pipeline/test_run_chunked_pipeline.py @@ -0,0 +1,37 @@ +from zipline.pipeline import Pipeline, run_chunked_pipeline +from zipline.pipeline.data import USEquityPricing +from zipline.pipeline.factors import Returns +from zipline.testing import ZiplineTestCase +from zipline.testing.fixtures import WithEquityPricingPipelineEngine + + +class ChunkedPipelineTestCase(WithEquityPricingPipelineEngine, + ZiplineTestCase): + + 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), + }, + ) + sessions = self.nyse_calendar.all_sessions + start_date = sessions[sessions.get_loc(self.START_DATE) + 2] + + pipeline_result = self.pipeline_engine.run_pipeline( + pipe, + start_date=start_date, + end_date=self.END_DATE, + ) + chunked_result = run_chunked_pipeline( + engine=self.pipeline_engine, + pipeline=pipe, + start_date=start_date, + end_date=self.END_DATE, + chunksize=22 + ) + self.assertTrue(chunked_result.equals(pipeline_result)) diff --git a/tests/utils/test_date_utils.py b/tests/utils/test_date_utils.py index ab43e11e..b9f3cd05 100644 --- a/tests/utils/test_date_utils.py +++ b/tests/utils/test_date_utils.py @@ -1,5 +1,4 @@ from pandas import Timestamp - from nose_parameterized import parameterized from zipline.testing import ZiplineTestCase diff --git a/zipline/pipeline/__init__.py b/zipline/pipeline/__init__.py index a169256b..16c429c2 100644 --- a/zipline/pipeline/__init__.py +++ b/zipline/pipeline/__init__.py @@ -1,5 +1,8 @@ from __future__ import print_function from zipline.assets import AssetFinder +from zipline.utils.calendars import get_calendar +from zipline.utils.date_utils import compute_date_range_chunks +from zipline.utils.pandas_utils import categorical_df_concat from .classifiers import Classifier, CustomClassifier from .engine import SimplePipelineEngine @@ -47,6 +50,39 @@ def engine_from_files(daily_bar_path, ) +def run_chunked_pipeline(engine, pipeline, start_date, end_date, chunksize): + """Run a pipeline to collect the results. + + Parameters + ---------- + engine : Engine + The pipeline engine. + 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 this is None, all the days + will be run at once. + + Returns + ------- + results : pd.DataFrame + The results for each output term in the pipeline. + """ + ranges = compute_date_range_chunks( + get_calendar('NYSE'), + start_date, + end_date, + chunksize, + ) + chunks = [engine.run_pipeline(pipeline, s, e) for s, e in ranges] + + return categorical_df_concat(chunks, inplace=True) + + __all__ = ( 'Classifier', 'CustomFactor', @@ -58,6 +94,7 @@ __all__ = ( 'Filter', 'Pipeline', 'SimplePipelineEngine', + 'run_chunked_pipeline', 'Term', 'TermGraph', ) diff --git a/zipline/utils/date_utils.py b/zipline/utils/date_utils.py index 07ceeda9..c59109a2 100644 --- a/zipline/utils/date_utils.py +++ b/zipline/utils/date_utils.py @@ -1,11 +1,15 @@ -def roll_dates_to_previous_session(calendar, *dates): +from toolz import partition_all + + +def roll_dates_to_previous_session(sessions, *dates): """ - Roll ``dates`` to the next session of ``calendar``. + Roll `dates` to the last session of `calendar`. Return input date if it + is a valid session. Parameters ---------- - calendar : zipline.utils.calendars.trading_calendar.TradingCalendar - The calendar to use as a reference. + sessions : pandas.tseries.index.DatetimeIndex + The list of valid session dates. *dates : pd.Timestamp The dates for which the last trading date is needed. @@ -15,7 +19,47 @@ def roll_dates_to_previous_session(calendar, *dates): The last trading date of the input dates, inclusive. """ - all_sessions = calendar.all_sessions + # Find the previous index value if there is no exact match. + locs = [sessions.get_loc(dt, method='ffill') for dt in dates] + return sessions[locs].tolist() - locs = [all_sessions.get_loc(dt, method='ffill') for dt in dates] - return all_sessions[locs] + +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] + ) + ) diff --git a/zipline/utils/pandas_utils.py b/zipline/utils/pandas_utils.py index 39e4040b..7323e245 100644 --- a/zipline/utils/pandas_utils.py +++ b/zipline/utils/pandas_utils.py @@ -258,7 +258,7 @@ def categorical_df_concat(df_list, inplace=False): new_categories = sorted( set().union( *(frame[col].cat.categories for frame in df_list) - ) + ) - {None} ) for df in df_list: