diff --git a/tests/test_restrictions.py b/tests/test_restrictions.py new file mode 100644 index 00000000..c1817b68 --- /dev/null +++ b/tests/test_restrictions.py @@ -0,0 +1,230 @@ +import pandas as pd +from pandas.util.testing import assert_series_equal +from nose_parameterized import parameterized +from six import iteritems +from functools import partial + +from zipline.finance.restrictions import ( + RESTRICTION_STATES, + Restriction, + HistoricalRestrictions, + StaticRestrictions, + NoopRestrictions, +) + +from zipline.testing.fixtures import ( + WithDataPortal, + ZiplineTestCase, +) + +str_to_ts = lambda dt_str: pd.Timestamp(dt_str, tz='UTC') +FROZEN = RESTRICTION_STATES.FROZEN +ALLOWED = RESTRICTION_STATES.ALLOWED +MINUTE = pd.Timedelta(minutes=1) + + +class RestrictionsTestCase(WithDataPortal, ZiplineTestCase): + + ASSET_FINDER_EQUITY_SIDS = 1, 2, 3 + + @classmethod + def init_class_fixtures(cls): + super(RestrictionsTestCase, cls).init_class_fixtures() + cls.ASSET1 = cls.asset_finder.retrieve_asset(1) + cls.ASSET2 = cls.asset_finder.retrieve_asset(2) + cls.ASSET3 = cls.asset_finder.retrieve_asset(3) + + def assert_is_restricted(self, rl, asset, dt): + self.assertTrue(rl.is_restricted(asset, dt)) + + def assert_not_restricted(self, rl, asset, dt): + self.assertFalse(rl.is_restricted(asset, dt)) + + def assert_vectorized_results(self, rl, expected, dt): + assert_series_equal( + rl.is_restricted([self.ASSET1, self.ASSET2, self.ASSET3], dt), + pd.Series( + index=pd.Index([self.ASSET1, self.ASSET2, self.ASSET3]), + data=expected + ) + ) + + @parameterized.expand([ + ('_'.join([timing, ordering]), + timing == 'intraday', + ordering == 'ordered') + for timing in ['intraday', 'interday'] + for ordering in ['ordered', 'unordered'] + ]) + def test_historical_restrictions(self, name, is_intraday, is_ordered): + """ + Test historical restrictions for both interday and intraday + restrictions, as well as restrictions defined in/not in order, for both + single- and multi-asset queries + """ + + hour_of_day = ' 15:00' if is_intraday else '' + + if is_ordered: + restriction_dates = { + self.ASSET1: [ + (str_to_ts('2011-01-04' + hour_of_day), FROZEN), + (str_to_ts('2011-01-05' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-06' + hour_of_day), FROZEN), + ], + self.ASSET2: [ + (str_to_ts('2011-01-05' + hour_of_day), FROZEN), + (str_to_ts('2011-01-06' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-07' + hour_of_day), FROZEN), + ], + } + else: + restriction_dates = { + self.ASSET1: [ + (str_to_ts('2011-01-05' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-06' + hour_of_day), FROZEN), + (str_to_ts('2011-01-04' + hour_of_day), FROZEN), + ], + self.ASSET2: [ + (str_to_ts('2011-01-06' + hour_of_day), ALLOWED), + (str_to_ts('2011-01-05' + hour_of_day), FROZEN), + (str_to_ts('2011-01-07' + hour_of_day), FROZEN), + ], + } + + restrictions = sum([ + [Restriction(asset, info[0], info[1]) for info in r_history] + for asset, r_history in iteritems(restriction_dates) + ], []) + rl = HistoricalRestrictions(restrictions) + + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + assert_vectorized_results = partial(self.assert_vectorized_results, rl) + + for asset, r_history in iteritems(restriction_dates): + dts = sorted([info[0] for info in r_history]) + + # Not restricted until on or after the freeze + assert_not_restricted(asset, dts[0] - MINUTE) + assert_is_restricted(asset, dts[0]) + assert_is_restricted(asset, dts[0] + MINUTE) + + # Unrestricted on or after the unfreeze + assert_is_restricted(asset, dts[1] - MINUTE) + assert_not_restricted(asset, dts[1]) + assert_not_restricted(asset, dts[1] + MINUTE) + + # Restricted again on or after the freeze + assert_not_restricted(asset, dts[2] - MINUTE) + assert_is_restricted(asset, dts[2]) + assert_is_restricted(asset, dts[2] + MINUTE) + # Should stay restricted for the rest of time + assert_is_restricted(asset, dts[2] + MINUTE * 1000000) + + dts = [str_to_ts(ts + hour_of_day) for ts in ['2011-01-04', + '2011-01-05', + '2011-01-06', + '2011-01-07']] + + # Expected results for [self.ASSET1, self.ASSET2, self.ASSET3], + # ASSET3 is always False as it has no defined restrictions + + # 01/04 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: ALLOWED + assert_vectorized_results([False, False, False], dts[0] - MINUTE) + assert_vectorized_results([True, False, False], dts[0]) + assert_vectorized_results([True, False, False], dts[0] + MINUTE) + + # 01/05 XX:00 ASSET1: FROZEN --> ALLOWED; ASSET2: ALLOWED --> FROZEN + assert_vectorized_results([True, False, False], dts[1] - MINUTE) + assert_vectorized_results([False, True, False], dts[1]) + assert_vectorized_results([False, True, False], dts[1] + MINUTE) + + # 01/06 XX:00 ASSET1: ALLOWED --> FROZEN; ASSET2: FROZEN --> ALLOWED + assert_vectorized_results([False, True, False], dts[2] - MINUTE) + assert_vectorized_results([True, False, False], dts[2]) + assert_vectorized_results([True, False, False], dts[2] + MINUTE) + + # 01/07 XX:00 ASSET1: FROZEN; ASSET2: ALLOWED --> FROZEN + assert_vectorized_results([True, False, False], dts[3] - MINUTE) + assert_vectorized_results([True, True, False], dts[3]) + assert_vectorized_results([True, True, False], dts[3] + MINUTE) + # Should stay restricted for the rest of time + assert_vectorized_results( + [True, True, False], + dts[3] + MINUTE * 10000000 + ) + + def test_historical_restrictions_consecutive_states(self): + """ + Test that defining redundant consecutive restrictions still works + """ + rl = HistoricalRestrictions([ + Restriction(self.ASSET1, str_to_ts('2011-01-04'), ALLOWED), + Restriction(self.ASSET1, str_to_ts('2011-01-05'), ALLOWED), + Restriction(self.ASSET1, str_to_ts('2011-01-06'), FROZEN), + Restriction(self.ASSET1, str_to_ts('2011-01-07'), FROZEN), + ]) + + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + + # (implicit) ALLOWED --> ALLOWED + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04') - MINUTE) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04')) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-04') + MINUTE) + + # ALLOWED --> ALLOWED + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05') - MINUTE) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05')) + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-05') + MINUTE) + + # ALLOWED --> FROZEN + assert_not_restricted(self.ASSET1, str_to_ts('2011-01-06') - MINUTE) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-06')) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-06') + MINUTE) + + # FROZEN --> FROZEN + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07') - MINUTE) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07')) + assert_is_restricted(self.ASSET1, str_to_ts('2011-01-07') + MINUTE) + + def test_static_restrictions(self): + """ + Test single- and multi-asset queries on static restrictions + """ + + restricted_a1 = self.ASSET1 + restricted_a2 = self.ASSET2 + unrestricted_a3 = self.ASSET3 + + rl = StaticRestrictions([restricted_a1, restricted_a2]) + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_is_restricted = partial(self.assert_is_restricted, rl) + assert_vectorized_results = partial(self.assert_vectorized_results, rl) + + for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03', + '2011-01-04', + '2020-01-04')]: + assert_is_restricted(restricted_a1, dt) + assert_is_restricted(restricted_a2, dt) + assert_not_restricted(unrestricted_a3, dt) + + assert_vectorized_results([True, True, False], dt) + + def test_noop_restrictions(self): + """ + Test single- and multi-asset queries on no-op restrictions + """ + + rl = NoopRestrictions() + assert_not_restricted = partial(self.assert_not_restricted, rl) + assert_vectorized_results = partial(self.assert_vectorized_results, rl) + + for dt in [str_to_ts(dt_str) for dt_str in ('2011-01-03', + '2011-01-04', + '2020-01-04')]: + assert_not_restricted(self.ASSET1, dt) + assert_not_restricted(self.ASSET2, dt) + assert_not_restricted(self.ASSET3, dt) + assert_vectorized_results([False, False, False], dt) diff --git a/zipline/finance/restrictions.py b/zipline/finance/restrictions.py new file mode 100644 index 00000000..e977d2d8 --- /dev/null +++ b/zipline/finance/restrictions.py @@ -0,0 +1,129 @@ +import abc +from numpy import vectorize +from functools import partial +import pandas as pd +from six import with_metaclass +from collections import namedtuple +from itertools import groupby + +from zipline.utils.enum import enum +from zipline.utils.numpy_utils import vectorized_is_element +from zipline.assets import Asset + + +Restriction = namedtuple( + 'Restriction', ['asset', 'effective_date', 'state'] +) + + +RESTRICTION_STATES = enum( + 'ALLOWED', + 'FROZEN', +) + + +class Restrictions(with_metaclass(abc.ABCMeta)): + """ + Abstract restricted list interface + """ + + @abc.abstractmethod + def is_restricted(self, assets, dt): + """ + Is the asset restricted (RestrictionStates.FROZEN) on the given dt? + + Parameters + ---------- + asset : Asset of iterable of Assets + The asset(s) for which we are querying a restriction + dt : pd.Timestamp + The timestamp of the restriction query + + Returns + ------- + is_restricted : bool or pd.Series[bool] indexed by asset + Is the asset or assets restricted on this dt? + + """ + raise NotImplementedError('is_restricted') + + +class NoopRestrictions(Restrictions): + """ + A no-op restrictions that contains no restrictions + """ + def is_restricted(self, assets, dt): + if isinstance(assets, Asset): + return False + return pd.Series(index=pd.Index(assets), data=[False]*len(assets)) + + +class StaticRestrictions(Restrictions): + """ + Static restrictions stored in memory that are constant regardless of dt + for each asset + + Parameters + ---------- + restricted_list : iterable of assets + The assets to be restricted + """ + + def __init__(self, restricted_list): + self._restricted_set = frozenset(restricted_list) + + def is_restricted(self, assets, dt): + """ + An asset is restricted for all dts if it is in the static list + """ + if isinstance(assets, Asset): + return assets in self._restricted_set + return pd.Series( + index=pd.Index(assets), + data=vectorized_is_element(assets, self._restricted_set) + ) + + +class HistoricalRestrictions(Restrictions): + """ + Historical restrictions stored in memory with effective dates for each + asset + + Parameters + ---------- + restrictions : iterable of namedtuple Restriction + The restrictions, each defined by an asset, effective date and state + """ + + def __init__(self, restrictions): + # A dict mapping each asset to its restrictions, which are sorted by + # ascending order of effective_date + self._restrictions_by_asset = { + asset: sorted( + restrictions_for_asset, key=lambda x: x.effective_date + ) + for asset, restrictions_for_asset + in groupby(restrictions, lambda x: x.asset) + } + + def is_restricted(self, assets, dt): + """ + Returns whether or not an asset or iterable of assets is restricted + on a dt + """ + if isinstance(assets, Asset): + return self._is_restricted_for_asset(assets, dt) + + is_restricted = partial(self._is_restricted_for_asset, dt=dt) + return pd.Series( + index=pd.Index(assets), + data=vectorize(is_restricted, otypes=[bool])(assets) + ) + + def _is_restricted_for_asset(self, asset, dt): + state = RESTRICTION_STATES.ALLOWED + for r in self._restrictions_by_asset.get(asset, ()): + if r.effective_date > dt: + break + state = r.state + return state == RESTRICTION_STATES.FROZEN