From ede1eb7aa021a3587da7b2a5bd2ba14daa763e58 Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Wed, 3 Feb 2016 16:09:14 -0500 Subject: [PATCH 1/3] PERF: Look up expired futures from in-memory Futures instead of queries to the db. --- zipline/assets/_assets.pyx | 8 ++++ zipline/assets/assets.py | 24 ---------- zipline/gens/tradesimulation.py | 30 ++++++++----- zipline/protocol.py | 2 +- zipline/utils/data.py | 80 ++++++++++++++++++++++++++++++++- 5 files changed, 108 insertions(+), 36 deletions(-) diff --git a/zipline/assets/_assets.pyx b/zipline/assets/_assets.pyx index 371253f8..e3cac229 100644 --- a/zipline/assets/_assets.pyx +++ b/zipline/assets/_assets.pyx @@ -230,6 +230,7 @@ cdef class Future(Asset): cdef readonly object auto_close_date cdef readonly object tick_size cdef readonly float multiplier + cdef readonly object effective_expiration def __cinit__(self, int sid, # sid is required @@ -253,6 +254,13 @@ cdef class Future(Asset): self.tick_size = tick_size self.multiplier = multiplier + if notice_date is None: + self.effective_expiration = expiration_date + elif expiration_date is None: + self.effective_expiration = notice_date + else: + self.effective_expiration = min(notice_date, expiration_date) + def __str__(self): if self.symbol: return 'Future(%d [%s])' % (self.sid, self.symbol) diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 3124cab0..1af5b571 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -666,30 +666,6 @@ class AssetFinder(object): contracts = self.retrieve_futures_contracts(sids) return [contracts[sid] for sid in sids] - def lookup_expired_futures(self, start, end): - if not isinstance(start, pd.Timestamp): - start = pd.Timestamp(start) - start = start.value - if not isinstance(end, pd.Timestamp): - end = pd.Timestamp(end) - end = end.value - - fc_cols = self.futures_contracts.c - - nd = sa.func.nullif(fc_cols.notice_date, pd.tslib.iNaT) - ed = sa.func.nullif(fc_cols.expiration_date, pd.tslib.iNaT) - date = sa.func.coalesce(sa.func.min(nd, ed), ed, nd) - - sids = list(map( - itemgetter('sid'), - sa.select((fc_cols.sid,)).where( - (date >= start) & (date < end)).order_by( - sa.func.coalesce(ed, nd).asc() - ).execute().fetchall() - )) - - return sids - @property def sids(self): return tuple(map( diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index f720333b..3a8459c1 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -12,20 +12,21 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from datetime import timedelta +from itertools import takewhile from contextlib2 import ExitStack - from logbook import Logger, Processor from pandas.tslib import normalize_date -from zipline.utils.api_support import ZiplineAPI - from zipline.finance.trading import NoFurtherDataError from zipline.protocol import ( BarData, SIDData, DATASOURCE_TYPE ) +from zipline.utils.api_support import ZiplineAPI +from zipline.utils.data import SortedDict log = Logger('Trade Simulation') @@ -56,15 +57,23 @@ class AlgorithmSimulator(object): # Snapshot Setup # ============== + def _get_effective_expiration(sid, + finder=self.env.asset_finder, + default=self.sim_params.last_close + + timedelta(days=1)): + asset = finder.retrieve_asset(sid) + return getattr(asset, 'effective_expiration', None) or default + + self._get_expiration = _get_effective_expiration + # The algorithm's data as of our most recent event. # We want an object that will have empty objects as default # values on missing keys. - self.current_data = BarData() + self.current_data = BarData(SortedDict(self._get_expiration)) # We don't have a datetime for the current snapshot until we # receive a message. self.simulation_dt = None - self.previous_dt = self.algo_start # ============= # Logging Setup @@ -97,14 +106,15 @@ class AlgorithmSimulator(object): self._call_before_trading_start(mkt_open) for date, snapshot in stream_in: - expired_sids = self.env.asset_finder.lookup_expired_futures( - start=self.previous_dt, end=date) - self.previous_dt = date + self.simulation_dt = date self.on_dt_changed(date) - # removing expired futures - for sid in expired_sids: + expired = list(takewhile( + lambda asset_id: self._get_expiration(asset_id) < date, + self.current_data + )) + for sid in expired: try: del self.current_data[sid] except KeyError: diff --git a/zipline/protocol.py b/zipline/protocol.py index cf4aca63..0355c54a 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -502,7 +502,7 @@ class BarData(object): """ def __init__(self, data=None): - self._data = data or {} + self._data = data if data is not None else {} self._contains_override = None def __contains__(self, name): diff --git a/zipline/utils/data.py b/zipline/utils/data.py index d81fd87d..1f4788b5 100644 --- a/zipline/utils/data.py +++ b/zipline/utils/data.py @@ -13,11 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import bisect import datetime +from collections import MutableMapping +from copy import deepcopy + +try: + from six.moves._thread import get_ident +except ImportError: + from six.moves._dummy_thread import get_ident import numpy as np import pandas as pd -from copy import deepcopy +from toolz import merge def _ensure_index(x): @@ -391,3 +399,73 @@ class MutableIndexRollingPanel(object): self.buffer.loc[non_nan_items, :, non_nan_cols]) self.buffer = new_buffer + + +class SortedDict(MutableMapping): + """A mapping of key-value pairs sorted by key according to the sort_key + function provided to the mapping. Ties from the sort_key are broken by + comparing the original keys. `iter` traverses the keys in sort order. + + Parameters + ---------- + key : callable + Called on keys in the mapping to produce the values by which those keys + are sorted. + mapping : mapping, optional + **kwargs + The initial mapping. + + >>> d = SortedDict(abs) + >>> d[-1] = 'negative one' + >>> d[0] = 'zero' + >>> d[2] = 'two' + >>> d + SortedDict[]([(0, 'zero'), (-1, 'negative one'), (1, 'one')]) # noqa + >>> d[1] = 'one' # Mutating the mapping maintains the sort order. + >>> d + SortedDict[]([(0, 'zero'), (-1, 'negative one'), (1, 'one')]) # noqa + """ + def __init__(self, sort_key, mapping=None, **kwargs): + self._map = {} + self._sorted_key_names = [] + self._sort_key = sort_key + + self.update(merge(mapping or {}, kwargs)) + + def __getitem__(self, name): + return self._map[name] + + def __setitem__(self, name, value, _bisect_right=bisect.bisect_right): + self._map[name] = value + if len(self._map) > len(self._sorted_key_names): + key = self._sort_key(name) + pair = (key, name) + idx = _bisect_right(self._sorted_key_names, pair) + self._sorted_key_names.insert(idx, pair) + + def __delitem__(self, name, _bisect_left=bisect.bisect_left): + del self._map[name] + idx = _bisect_left(self._sorted_key_names, + (self._sort_key(name), name)) + del self._sorted_key_names[idx] + + def __iter__(self): + for key, name in self._sorted_key_names: + yield name + + def __len__(self): + return len(self._map) + + def __repr__(self, _repr_running={}): + # Based on OrderedDict/defaultdict + call_key = id(self), get_ident() + if call_key in _repr_running: + return '...' + _repr_running[call_key] = 1 + try: + if not self: + return '%s[%r]()' % (self.__class__.__name__, self._sort_key) + return '%s[%r](%r)' % (self.__class__.__name__, self._sort_key, + self.items()) + finally: + del _repr_running[call_key] From 79aff84aed3fdad1789a5695de37f950a9b5ca15 Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Thu, 4 Feb 2016 10:58:08 -0500 Subject: [PATCH 2/3] MAINT: Moved auto_close_date to Asset Handle custom data source with non-int sids --- zipline/assets/_assets.pyx | 28 ++++++++++++----------- zipline/gens/tradesimulation.py | 40 ++++++++++++++++++++++----------- 2 files changed, 42 insertions(+), 26 deletions(-) diff --git a/zipline/assets/_assets.pyx b/zipline/assets/_assets.pyx index e3cac229..367dcbf1 100644 --- a/zipline/assets/_assets.pyx +++ b/zipline/assets/_assets.pyx @@ -37,7 +37,7 @@ cimport numpy as np # IMPORTANT NOTE: You must change this template if you change # Asset.__reduce__, or else we'll attempt to unpickle an old version of this # class -CACHE_FILE_TEMPLATE = '/tmp/.%s-%s.v5.cache' +CACHE_FILE_TEMPLATE = '/tmp/.%s-%s.v6.cache' cdef class Asset: @@ -51,6 +51,7 @@ cdef class Asset: cdef readonly object start_date cdef readonly object end_date cdef public object first_traded + cdef readonly object auto_close_date cdef readonly object exchange @@ -61,6 +62,7 @@ cdef class Asset: object start_date=None, object end_date=None, object first_traded=None, + object auto_close_date=None, object exchange="", *args, **kwargs): @@ -73,6 +75,7 @@ cdef class Asset: self.start_date = start_date self.end_date = end_date self.first_traded = first_traded + self.auto_close_date = auto_close_date def __int__(self): return self.sid @@ -127,7 +130,7 @@ cdef class Asset: def __repr__(self): attrs = ('symbol', 'asset_name', 'exchange', - 'start_date', 'end_date', 'first_traded') + 'start_date', 'end_date', 'first_traded', 'auto_close_date') tuples = ((attr, repr(getattr(self, attr, None))) for attr in attrs) strings = ('%s=%s' % (t[0], t[1]) for t in tuples) @@ -147,6 +150,7 @@ cdef class Asset: self.start_date, self.end_date, self.first_traded, + self.auto_close_date, self.exchange,)) cpdef to_dict(self): @@ -160,6 +164,7 @@ cdef class Asset: 'start_date': self.start_date, 'end_date': self.end_date, 'first_traded': self.first_traded, + 'auto_close_date': self.auto_close_date, 'exchange': self.exchange, } @@ -181,7 +186,7 @@ cdef class Equity(Asset): def __repr__(self): attrs = ('symbol', 'asset_name', 'exchange', - 'start_date', 'end_date', 'first_traded') + 'start_date', 'end_date', 'first_traded', 'auto_close_date') tuples = ((attr, repr(getattr(self, attr, None))) for attr in attrs) strings = ('%s=%s' % (t[0], t[1]) for t in tuples) @@ -227,10 +232,8 @@ cdef class Future(Asset): cdef readonly object root_symbol cdef readonly object notice_date cdef readonly object expiration_date - cdef readonly object auto_close_date cdef readonly object tick_size cdef readonly float multiplier - cdef readonly object effective_expiration def __cinit__(self, int sid, # sid is required @@ -250,16 +253,16 @@ cdef class Future(Asset): self.root_symbol = root_symbol self.notice_date = notice_date self.expiration_date = expiration_date - self.auto_close_date = auto_close_date self.tick_size = tick_size self.multiplier = multiplier - if notice_date is None: - self.effective_expiration = expiration_date - elif expiration_date is None: - self.effective_expiration = notice_date - else: - self.effective_expiration = min(notice_date, expiration_date) + if auto_close_date is None: + if notice_date is None: + self.auto_close_date = expiration_date + elif expiration_date is None: + self.auto_close_date = notice_date + else: + self.auto_close_date = min(notice_date, expiration_date) def __str__(self): if self.symbol: @@ -307,7 +310,6 @@ cdef class Future(Asset): super_dict['root_symbol'] = self.root_symbol super_dict['notice_date'] = self.notice_date super_dict['expiration_date'] = self.expiration_date - super_dict['auto_close_date'] = self.auto_close_date super_dict['tick_size'] = self.tick_size super_dict['multiplier'] = self.multiplier return super_dict diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 3a8459c1..23ae30db 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -19,6 +19,7 @@ from contextlib2 import ExitStack from logbook import Logger, Processor from pandas.tslib import normalize_date +from zipline.errors import SidsNotFound from zipline.finance.trading import NoFurtherDataError from zipline.protocol import ( BarData, @@ -57,19 +58,32 @@ class AlgorithmSimulator(object): # Snapshot Setup # ============== - def _get_effective_expiration(sid, - finder=self.env.asset_finder, - default=self.sim_params.last_close - + timedelta(days=1)): - asset = finder.retrieve_asset(sid) - return getattr(asset, 'effective_expiration', None) or default + def _get_asset_close_date(sid, + finder=self.env.asset_finder, + default=self.sim_params.last_close + + timedelta(days=1)): + try: + asset = finder.retrieve_asset(sid) + except ValueError: + # Handle sid not an int, such as from a custom source. + # So that they don't compare equal to other sids, and we'd + # blow up comparing strings to ints, let's give them unique + # close dates. + return default + timedelta(microseconds=id(sid)) + except SidsNotFound: + return default + # Default is used when the asset has no auto close date, + # and is set to a time after the simulation ends, so that the + # relevant asset isn't removed from the universe at all + # (at least not for this reason). + return asset.auto_close_date or default - self._get_expiration = _get_effective_expiration + self._get_asset_close = _get_asset_close_date # The algorithm's data as of our most recent event. - # We want an object that will have empty objects as default - # values on missing keys. - self.current_data = BarData(SortedDict(self._get_expiration)) + # Maintain sids in order by asset close date, so that we can more + # efficiently remove them when their times come... + self.current_data = BarData(SortedDict(self._get_asset_close)) # We don't have a datetime for the current snapshot until we # receive a message. @@ -110,11 +124,11 @@ class AlgorithmSimulator(object): self.simulation_dt = date self.on_dt_changed(date) - expired = list(takewhile( - lambda asset_id: self._get_expiration(asset_id) < date, + closed = list(takewhile( + lambda asset_id: self._get_asset_close(asset_id) < date, self.current_data )) - for sid in expired: + for sid in closed: try: del self.current_data[sid] except KeyError: From 4e84b2c5ca442527e889a805f7dad5a9cad85402 Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Thu, 4 Feb 2016 17:36:40 -0500 Subject: [PATCH 3/3] TST: Added doctest --- tests/test_doctests.py | 4 ++++ zipline/utils/data.py | 25 +++++++++++++++++-------- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/tests/test_doctests.py b/tests/test_doctests.py index 7ff2d5a2..c08659af 100644 --- a/tests/test_doctests.py +++ b/tests/test_doctests.py @@ -10,6 +10,7 @@ from zipline.pipeline import ( ) from zipline.utils import ( cache, + data, input_validation, memoize, numpy_utils, @@ -78,3 +79,6 @@ class DoctestTestCase(TestCase): def test_numpy_utils_docs(self): self._check_docs(numpy_utils) + + def test_data_docs(self): + self._check_docs(data) diff --git a/zipline/utils/data.py b/zipline/utils/data.py index 1f4788b5..9c10b8d7 100644 --- a/zipline/utils/data.py +++ b/zipline/utils/data.py @@ -419,16 +419,25 @@ class SortedDict(MutableMapping): >>> d[-1] = 'negative one' >>> d[0] = 'zero' >>> d[2] = 'two' - >>> d - SortedDict[]([(0, 'zero'), (-1, 'negative one'), (1, 'one')]) # noqa + >>> d # doctest: +NORMALIZE_WHITESPACE + SortedDict(, + [(0, 'zero'), (-1, 'negative one'), (2, 'two')]) >>> d[1] = 'one' # Mutating the mapping maintains the sort order. + >>> d # doctest: +NORMALIZE_WHITESPACE + SortedDict(, + [(0, 'zero'), (-1, 'negative one'), (1, 'one'), (2, 'two')]) + >>> del d[0] + >>> d # doctest: +NORMALIZE_WHITESPACE + SortedDict(, + [(-1, 'negative one'), (1, 'one'), (2, 'two')]) + >>> del d[2] >>> d - SortedDict[]([(0, 'zero'), (-1, 'negative one'), (1, 'one')]) # noqa + SortedDict(, [(-1, 'negative one'), (1, 'one')]) """ - def __init__(self, sort_key, mapping=None, **kwargs): + def __init__(self, key, mapping=None, **kwargs): self._map = {} self._sorted_key_names = [] - self._sort_key = sort_key + self._sort_key = key self.update(merge(mapping or {}, kwargs)) @@ -464,8 +473,8 @@ class SortedDict(MutableMapping): _repr_running[call_key] = 1 try: if not self: - return '%s[%r]()' % (self.__class__.__name__, self._sort_key) - return '%s[%r](%r)' % (self.__class__.__name__, self._sort_key, - self.items()) + return '%s(%r)' % (self.__class__.__name__, self._sort_key) + return '%s(%r, %r)' % (self.__class__.__name__, self._sort_key, + list(self.items())) finally: del _repr_running[call_key]