Merge pull request #983 from quantopian/expiration_perf

Expiration perf

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)

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,7 +232,6 @@ 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
 
@@ -249,10 +253,17 @@ 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 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:
             return 'Future(%d [%s])' % (self.sid, self.symbol)
@@ -299,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

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(

zipline/gens/tradesimulation.py

@@ -12,20 +12,22 @@
 # 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.errors import SidsNotFound
 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 +58,36 @@ class AlgorithmSimulator(object):
         # Snapshot Setup
         # ==============
 
+        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_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()
+        # 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.
         self.simulation_dt = None
-        self.previous_dt = self.algo_start
 
         # =============
         # Logging Setup
@@ -97,14 +120,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:
+                closed = list(takewhile(
+                    lambda asset_id: self._get_asset_close(asset_id) < date,
+                    self.current_data
+                ))
+                for sid in closed:
                     try:
                         del self.current_data[sid]
                     except KeyError:

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):

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,82 @@ 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  # doctest: +NORMALIZE_WHITESPACE
+    SortedDict(<built-in function abs>,
+               [(0, 'zero'), (-1, 'negative one'), (2, 'two')])
+    >>> d[1] = 'one'  # Mutating the mapping maintains the sort order.
+    >>> d  # doctest: +NORMALIZE_WHITESPACE
+    SortedDict(<built-in function abs>,
+               [(0, 'zero'), (-1, 'negative one'), (1, 'one'), (2, 'two')])
+    >>> del d[0]
+    >>> d  # doctest: +NORMALIZE_WHITESPACE
+    SortedDict(<built-in function abs>,
+               [(-1, 'negative one'), (1, 'one'), (2, 'two')])
+    >>> del d[2]
+    >>> d
+    SortedDict(<built-in function abs>, [(-1, 'negative one'), (1, 'one')])
+    """
+    def __init__(self, key, mapping=None, **kwargs):
+        self._map = {}
+        self._sorted_key_names = []
+        self._sort_key = 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,
+                                   list(self.items()))
+        finally:
+            del _repr_running[call_key]