Merge pull request #1529 from quantopian/current-contract

ENH: Add continuous future current contract.

zipline/_protocol.pyx

@@ -25,6 +25,7 @@ from cpython cimport bool
 from collections import Iterable
 
 from zipline.assets import Asset, Future
+from zipline.assets.continuous_futures import ContinuousFuture
 from zipline.zipline_warnings import ZiplineDeprecationWarning
 
 
@@ -254,7 +255,7 @@ cdef class BarData:
         return dt
 
     @check_parameters(('assets', 'fields'),
-                      ((Asset,) + string_types, string_types))
+                      ((Asset, ContinuousFuture) + string_types, string_types))
     def current(self, assets, fields):
         """
         Returns the current value of the given assets for the given fields

zipline/algorithm.py

@@ -1185,6 +1185,33 @@ class TradingAlgorithm(object):
             as_of_date=_lookup_date,
         )
 
+    @api_method
+    @preprocess(root_symbol_str=ensure_upper_case)
+    def continuous_future(self, root_symbol_str, offset, roll):
+        """Create a specifier for a continuous contract.
+
+        Parameters
+        ----------
+        root_symbol_str : str
+            The root symbol for the future chain.
+
+        offset : int
+            The distance from the primary contract.
+
+        roll_style : str
+            How rolls are determined.
+
+        Returns
+        -------
+        continuous_future : ContinuousFuture
+            The continuous future specifier.
+        """
+        return self.asset_finder.create_continuous_future(
+            root_symbol_str,
+            offset,
+            roll,
+        )
+
     @api_method
     def symbols(self, *args):
         """Lookup multuple Equities as a list.

zipline/assets/assets.py

@@ -13,9 +13,12 @@
 # limitations under the License.
 
 from abc import ABCMeta
+import array
+import binascii
 from collections import namedtuple
 from numbers import Integral
 from operator import itemgetter, attrgetter
+import struct
 
 from logbook import Logger
 import numpy as np
@@ -46,6 +49,7 @@ from zipline.errors import (
 from . import (
     Asset, Equity, Future,
 )
+from . continuous_futures import OrderedContracts, ContinuousFuture
 from .asset_writer import (
     check_version_info,
     split_delimited_symbol,
@@ -119,6 +123,41 @@ def _convert_asset_timestamp_fields(dict_):
     return dict_
 
 
+SID_TYPE_IDS = {
+    # Asset would be 0,
+    ContinuousFuture: 1,
+}
+
+CONTINUOUS_FUTURE_ROLL_STYLE_IDS = {
+    'calendar': 0
+}
+
+
+def _encode_continuous_future_sid(root_symbol, offset, roll_style):
+    s = struct.Struct("B 2B B B 3B")
+    # B - sid type
+    # 2B - root symbol
+    # B - offset (could be packed smaller since offsets of greater than 12 are
+    #             probably unneeded.)
+    # B - roll type
+    # 3B - empty space left for parameterized roll types
+
+    # The root symbol currently supports 2 characters.  If 3 char root symbols
+    # are needed, the size of the root symbol does not need to change, however
+    # writing the string directly will need to change to a scheme of writing
+    # the A-Z values in 5-bit chunks.
+    a = array.array('B', [0] * s.size)
+    rs = bytearray(root_symbol, 'ascii')
+    values = (SID_TYPE_IDS[ContinuousFuture],
+              rs[0],
+              rs[1],
+              offset,
+              CONTINUOUS_FUTURE_ROLL_STYLE_IDS[roll_style],
+              0, 0, 0,)
+    s.pack_into(a, 0, *values)
+    return int(binascii.hexlify(a), 16)
+
+
 class AssetFinder(object):
     """
     An AssetFinder is an interface to a database of Asset metadata written by
@@ -162,6 +201,8 @@ class AssetFinder(object):
         # retrieve_asset will populate the cache on first retrieval.
         self._caches = (self._asset_cache, self._asset_type_cache) = {}, {}
 
+        self._ordered_contracts = {}
+
         # Populated on first call to `lifetimes`.
         self._asset_lifetimes = None
 
@@ -821,6 +862,63 @@ class AssetFinder(object):
 
         return sids
 
+    def _get_contract_info(self, root_symbol):
+        fc_cols = self.futures_contracts.c
+
+        fields = (fc_cols.sid, fc_cols.start_date, fc_cols.auto_close_date)
+
+        return list(sa.select(fields).where(
+            (fc_cols.root_symbol == root_symbol) &
+            (fc_cols.start_date != pd.NaT.value))
+            .order_by(fc_cols.auto_close_date).execute().fetchall())
+
+    def _get_root_symbol_exchange(self, root_symbol):
+        fc_cols = self.futures_root_symbols.c
+
+        fields = (fc_cols.exchange,)
+
+        return sa.select(fields).where(
+            fc_cols.root_symbol == root_symbol).execute().fetchone()[0]
+
+    def get_ordered_contracts(self, root_symbol):
+        try:
+            return self._ordered_contracts[root_symbol]
+        except KeyError:
+            contract_info = self._get_contract_info(root_symbol)
+            size = len(contract_info)
+            sids = np.full(size, 0, dtype=np.int64)
+            start_dates = np.full(size, 0, dtype=np.int64)
+            auto_close_dates = np.full(size, 0, dtype=np.int64)
+            self._size = size
+            for i, info in enumerate(contract_info):
+                sid, start_date, auto_close_date = info
+                sids[i] = sid
+                start_dates[i] = start_date
+                auto_close_dates[i] = auto_close_date
+            oc = OrderedContracts(root_symbol,
+                                  sids,
+                                  start_dates,
+                                  auto_close_dates)
+            self._ordered_contracts[root_symbol] = oc
+            return oc
+
+    def create_continuous_future(self, root_symbol, offset, roll_style):
+        oc = self.get_ordered_contracts(root_symbol)
+        start_date = self.retrieve_asset(oc.contract_sids[0]).start_date
+        end_date = self.retrieve_asset(oc.contract_sids[-1]).end_date
+        exchange = self._get_root_symbol_exchange(root_symbol)
+
+        sid = _encode_continuous_future_sid(root_symbol, offset, roll_style)
+        cf = ContinuousFuture(sid,
+                              root_symbol,
+                              offset,
+                              roll_style,
+                              start_date,
+                              end_date,
+                              exchange)
+        self._asset_cache[cf.sid] = cf
+        return cf
+
     def _make_sids(tblattr):
         def _(self):
             return tuple(map(

zipline/assets/continuous_futures.pyx

@@ -0,0 +1,296 @@
+# cython: embedsignature=True
+#
+# Copyright 2016 Quantopian, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""
+Cythonized ContinuousFutures object.
+"""
+cimport cython
+from cpython.number cimport PyNumber_Index
+from cpython.object cimport (
+    Py_EQ,
+    Py_NE,
+    Py_GE,
+    Py_LE,
+    Py_GT,
+    Py_LT,
+)
+from cpython cimport bool
+
+import numpy as np
+from numpy cimport long_t, int64_t
+import warnings
+cimport numpy as np
+
+from zipline.utils.calendars import get_calendar
+
+
+cdef class ContinuousFuture:
+    """
+    Represents a specifier for a chain of future contracts, where the
+    coordinates for the chain are:
+    root_symbol : str
+        The root symbol of the contracts.
+    offset : int
+        The distance from the primary chain.
+        e.g. 0 specifies the primary chain, 1 the secondary, etc.
+    roll_style : str
+        How rolls from contract to contract should be calculated.
+        Currently supports 'calendar'.
+
+    Instances of this class are exposed to the algorithm.
+    """
+
+    cdef readonly long_t sid
+    # Cached hash of self.sid
+    cdef long_t sid_hash
+
+    cdef readonly object root_symbol
+    cdef readonly int offset
+    cdef readonly object roll_style
+
+    cdef readonly object start_date
+    cdef readonly object end_date
+
+    cdef readonly object exchange
+
+    _kwargnames = frozenset({
+        'sid',
+        'root_symbol',
+        'offset',
+        'start_date',
+        'end_date',
+        'exchange',
+    })
+
+    def __init__(self,
+                 long_t sid, # sid is required
+                 object root_symbol,
+                 int offset,
+                 object roll_style,
+                 object start_date,
+                 object end_date,
+                 object exchange):
+
+        self.sid = sid
+        self.sid_hash = hash(sid)
+        self.root_symbol = root_symbol
+        self.roll_style = roll_style
+        self.offset = offset
+        self.exchange = exchange
+        self.start_date = start_date
+        self.end_date = end_date
+
+    def __int__(self):
+        return self.sid
+
+    def __index__(self):
+        return self.sid
+
+    def __hash__(self):
+        return self.sid_hash
+
+    def __richcmp__(x, y, int op):
+        """
+        Cython rich comparison method.  This is used in place of various
+        equality checkers in pure python.
+        """
+        cdef int x_as_int, y_as_int
+
+        try:
+            x_as_int = PyNumber_Index(x)
+        except (TypeError, OverflowError):
+            return NotImplemented
+
+        try:
+            y_as_int = PyNumber_Index(y)
+        except (TypeError, OverflowError):
+            return NotImplemented
+
+        compared = x_as_int - y_as_int
+
+        # Handle == and != first because they're significantly more common
+        # operations.
+        if op == Py_EQ:
+            return compared == 0
+        elif op == Py_NE:
+            return compared != 0
+        elif op == Py_LT:
+            return compared < 0
+        elif op == Py_LE:
+            return compared <= 0
+        elif op == Py_GT:
+            return compared > 0
+        elif op == Py_GE:
+            return compared >= 0
+        else:
+            raise AssertionError('%d is not an operator' % op)
+
+    def __str__(self):
+        return '%s(%d [%s, %s, %s])' % (
+            type(self).__name__,
+            self.sid,
+            self.root_symbol,
+            self.offset,
+            self.roll_style
+        )
+
+    def __repr__(self):
+        attrs = ('root_symbol', 'offset', 'roll_style')
+        tuples = ((attr, repr(getattr(self, attr, None)))
+                  for attr in attrs)
+        strings = ('%s=%s' % (t[0], t[1]) for t in tuples)
+        params = ', '.join(strings)
+        return 'ContinuousFuture(%d, %s)' % (self.sid, params)
+
+    cpdef __reduce__(self):
+        """
+        Function used by pickle to determine how to serialize/deserialize this
+        class.  Should return a tuple whose first element is self.__class__,
+        and whose second element is a tuple of all the attributes that should
+        be serialized/deserialized during pickling.
+        """
+        return (self.__class__, (self.sid,
+                                 self.root_symbol,
+                                 self.start_date,
+                                 self.end_date,
+                                 self.offset,
+                                 self.roll_style,
+                                 self.exchange))
+
+    cpdef to_dict(self):
+        """
+        Convert to a python dict.
+        """
+        return {
+            'sid': self.sid,
+            'root_symbol': self.root_symbol,
+            'start_date': self.start_date,
+            'end_date': self.end_date,
+            'offset': self.offset,
+            'roll_style': self.roll_style,
+            'exchange': self.exchange,
+        }
+
+    @classmethod
+    def from_dict(cls, dict_):
+        """
+        Build an ContinuousFuture instance from a dict.
+        """
+        return cls(**dict_)
+
+    def is_alive_for_session(self, session_label):
+        """
+        Returns whether the continuous future is alive at the given dt.
+
+        Parameters
+        ----------
+        session_label: pd.Timestamp
+            The desired session label to check. (midnight UTC)
+
+        Returns
+        -------
+        boolean: whether the continuous is alive at the given dt.
+        """
+        cdef int64_t ref_start
+        cdef int64_t ref_end
+
+        ref_start = self.start_date.value
+        ref_end = self.end_date.value
+
+        return ref_start <= session_label.value <= ref_end
+
+    def is_exchange_open(self, dt_minute):
+        """
+        Parameters
+        ----------
+        dt_minute: pd.Timestamp (UTC, tz-aware)
+            The minute to check.
+
+        Returns
+        -------
+        boolean: whether the continuous futures's exchange is open at the
+        given minute.
+        """
+        calendar = get_calendar(self.exchange)
+        return calendar.is_open_on_minute(dt_minute)
+
+
+cdef class OrderedContracts(object):
+    """
+    A container for aligned values of a future contract chain, in sorted order
+    of their occurrence.
+    Used to get answers about contracts in relation to their auto close
+    dates and start dates.
+
+    The number of contracts for a given root symbol is ~250,
+    which is why search by comparison over the range of contracts is
+    used. At this size, this is faster than using an index or np.searchsorted.
+
+    Members
+    -------
+    root_symbol : str
+        The root symbol of the future contract chain.
+    contract_sids : long[:]
+        The contract sids in sorted order of occurrence.
+    start_dates : long[:]
+        The start dates of the contracts in the chain.
+        Corresponds by index with contract_sids.
+    auto_close_dates : long[:]
+        The auto close dates of the contracts in the chain.
+        Corresponds by index with contract_sids.
+
+    Instances of this class are used by the simulation engine, but not
+    exposed to the algorithm.
+    """
+
+    cdef readonly object root_symbol
+    cdef int _size
+    cdef readonly long_t[:] contract_sids
+    cdef readonly long_t[:] start_dates
+    cdef readonly long_t[:] auto_close_dates
+
+    def __init__(self,
+                 object root_symbol,
+                 long_t[:] contract_sids,
+                 long_t[:] start_dates,
+                 long_t[:] auto_close_dates):
+        self._size = len(contract_sids)
+        self.root_symbol = root_symbol
+        self.contract_sids = contract_sids
+        self.start_dates = start_dates
+        self.auto_close_dates = auto_close_dates
+    
+    cpdef long contract_before_auto_close(self, long_t dt_value):
+        """
+        Get the contract with next upcoming auto close date.
+        """
+        cdef Py_ssize_t i, auto_close_date
+        for i, auto_close_date in enumerate(self.auto_close_dates):
+            if auto_close_date > dt_value:
+                break
+        return self.contract_sids[i]
+
+    cpdef long contract_at_offset(self, long_t sid, Py_ssize_t offset):
+        """
+        Get the sid which is the given sid plus the offset distance.
+        An offset of 0 should be reflexive.
+        """
+        cdef Py_ssize_t i
+        cdef long_t[:] sids
+        sids = self.contract_sids
+        for i in range(self._size):
+            if sid == sids[i]:
+                return sids[i + offset]

zipline/assets/roll_finder.py

@@ -0,0 +1,63 @@
+#
+# Copyright 2016 Quantopian, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 abc import ABCMeta, abstractmethod
+from six import with_metaclass
+
+
+class RollFinder(with_metaclass(ABCMeta, object)):
+    """
+    Abstract base class for calculating when futures contracts are the active
+    contract.
+    """
+
+    @abstractmethod
+    def get_contract_center(self, root_symbol, dt, offset):
+        """
+        Parameters
+        ----------
+        root_symbol : str
+            The root symbol for the contract chain.
+        dt : Timestamp
+            The datetime for which to retrieve the current contract.
+        offset : int
+            The offset from the primary contract.
+            0 is the primary, 1 is the secondary, etc.
+
+        Returns
+        -------
+        Future
+            The active future contract at the given dt.
+        """
+        raise NotImplemented
+
+
+class CalendarRollFinder(RollFinder):
+    """
+    The CalendarRollFinder calculates contract rolls based purely on the
+    contract's auto close date.
+    """
+
+    def __init__(self, trading_calendar, asset_finder):
+        self.trading_calendar = trading_calendar
+        self.asset_finder = asset_finder
+
+    def get_contract_center(self, root_symbol, dt, offset):
+        oc = self.asset_finder.get_ordered_contracts(root_symbol)
+        session = self.trading_calendar.minute_to_session_label(dt)
+        primary_candidate = oc.contract_before_auto_close(session.value)
+
+        # Here is where a volume check would be.
+        primary = primary_candidate
+        return oc.contract_at_offset(primary, offset)

zipline/data/data_portal.py

@@ -23,6 +23,8 @@ from six import iteritems
 from six.moves import reduce
 
 from zipline.assets import Asset, Future, Equity
+from zipline.assets.continuous_futures import ContinuousFuture
+from zipline.assets.roll_finder import CalendarRollFinder
 from zipline.data.dispatch_bar_reader import (
     AssetDispatchMinuteBarReader,
     AssetDispatchSessionBarReader
@@ -54,7 +56,14 @@ from zipline.errors import (
 log = Logger('DataPortal')
 
 BASE_FIELDS = frozenset([
-    "open", "high", "low", "close", "volume", "price", "last_traded"
+    "open",
+    "high",
+    "low",
+    "close",
+    "volume",
+    "price",
+    "contract",
+    "last_traded",
 ])
 
 OHLCV_FIELDS = frozenset([
@@ -149,6 +158,11 @@ class DataPortal(object):
         else:
             self._last_trading_session = None
 
+        self._roll_finders = {
+            'calendar': CalendarRollFinder(self.trading_calendar,
+                                           self.asset_finder)
+        }
+
         aligned_equity_minute_reader = self._ensure_reader_aligned(
             equity_minute_reader)
         aligned_equity_session_reader = self._ensure_reader_aligned(
@@ -350,7 +364,8 @@ class DataPortal(object):
         # at it if it's on something like palladium and not AAPL (since our
         # own price data always wins when dealing with assets).
 
-        return not (field in BASE_FIELDS and isinstance(asset, Asset))
+        return not (field in BASE_FIELDS and
+                    (isinstance(asset, (Asset, ContinuousFuture))))
 
     def _get_fetcher_value(self, asset, field, dt):
         day = normalize_date(dt)
@@ -405,6 +420,8 @@ class DataPortal(object):
                 return 0
             elif field != "last_traded":
                 return np.NaN
+            elif field == "contract":
+                return None
 
         if data_frequency == "daily":
             return self._get_daily_data(asset, field, session_label)
@@ -414,6 +431,8 @@ class DataPortal(object):
             elif field == "price":
                 return self._get_minute_spot_value(asset, "close", dt,
                                                    ffill=True)
+            elif field == "contract":
+                return self._get_current_contract(asset, dt)
             else:
                 return self._get_minute_spot_value(asset, field, dt)
 
@@ -1219,3 +1238,10 @@ class DataPortal(object):
                 ret = np.nan
 
             return ret
+
+    def _get_current_contract(self, continuous_future, dt):
+        rf = self._roll_finders[continuous_future.roll_style]
+        return self.asset_finder.retrieve_asset(
+            rf.get_contract_center(continuous_future.root_symbol,
+                                   dt,
+                                   continuous_future.offset))