diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 568b5442..b1e38920 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -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)) diff --git a/tests/utils/test_date_utils.py b/tests/utils/test_date_utils.py new file mode 100644 index 00000000..67ac141a --- /dev/null +++ b/tests/utils/test_date_utils.py @@ -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." + ) diff --git a/tests/utils/test_pandas_utils.py b/tests/utils/test_pandas_utils.py index e84ed3fe..68d1efb8 100644 --- a/tests/utils/test_pandas_utils.py +++ b/tests/utils/test_pandas_utils.py @@ -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." + ) diff --git a/tests/utils/test_sharedoc.py b/tests/utils/test_sharedoc.py new file mode 100644 index 00000000..bed1cef9 --- /dev/null +++ b/tests/utils/test_sharedoc.py @@ -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__ + ) diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py index 1d28a834..00906f90 100644 --- a/zipline/pipeline/engine.py +++ b/zipline/pipeline/engine.py @@ -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 diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index f1439d56..11bf5ace 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -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, ) diff --git a/zipline/utils/calendars/__init__.py b/zipline/utils/calendars/__init__.py index d3f3ecf0..61c22e6b 100644 --- a/zipline/utils/calendars/__init__.py +++ b/zipline/utils/calendars/__init__.py @@ -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', ] diff --git a/zipline/utils/date_utils.py b/zipline/utils/date_utils.py new file mode 100644 index 00000000..22e57383 --- /dev/null +++ b/zipline/utils/date_utils.py @@ -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] + ) + ) diff --git a/zipline/utils/pandas_utils.py b/zipline/utils/pandas_utils.py index ccac273a..2cae7f90 100644 --- a/zipline/utils/pandas_utils.py +++ b/zipline/utils/pandas_utils.py @@ -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) diff --git a/zipline/utils/sharedoc.py b/zipline/utils/sharedoc.py index d0d7248c..6ba53607 100644 --- a/zipline/utils/sharedoc.py +++ b/zipline/utils/sharedoc.py @@ -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