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

ENH: A point-in-time restricted list with restrictions stored in memory

Browse Source 
An ABC Restrictions defines a group of restrictions responsible
for returning restriction information for sids on certain dts. An
InMemoryRestrictions is a point-in-time group of such restrictions,
with all restrictions and their dates passed in upon instantiation.
A StaticRestrictedList takes a list of sids, restricting them at all
dates
This commit is contained in:
Andrew Liang
2016-09-14 14:01:23 -04:00
parent 60b736b594
commit 0f47cf313d
2 changed files with 359 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_restrictions.py
+230
View File
@@ -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)
zipline/finance/restrictions.py
+129
View File
@@ -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