diff --git a/tests/test_continuous_futures.py b/tests/test_continuous_futures.py index b8dba528..6d8ea811 100644 --- a/tests/test_continuous_futures.py +++ b/tests/test_continuous_futures.py @@ -15,10 +15,12 @@ from textwrap import dedent +from numpy import array, int64 import pandas as pd from pandas import Timestamp, DataFrame from zipline import TradingAlgorithm +from zipline.assets.continuous_futures import OrderedContracts from zipline.testing.fixtures import ( WithCreateBarData, WithSimParams, @@ -49,27 +51,37 @@ class ContinuousFuturesTestCase(WithCreateBarData, @classmethod def make_futures_info(self): return DataFrame({ - 'symbol': ['FOF', 'FOG', 'FOH'], - 'root_symbol': ['FO', 'FO', 'FO'], - 'asset_name': ['Foo'] * 3, + 'symbol': ['FOF16', 'FOG16', 'FOH16', 'FOJ16', 'FOF22'], + 'root_symbol': ['FO', 'FO', 'FO', 'FO', 'FO'], + 'asset_name': ['Foo'] * 5, 'start_date': [Timestamp('2015-01-05', tz='UTC'), Timestamp('2015-02-05', tz='UTC'), - Timestamp('2015-03-05', tz='UTC')], + Timestamp('2015-03-05', tz='UTC'), + Timestamp('2015-04-05', tz='UTC'), + Timestamp('2021-01-05', tz='UTC')], 'end_date': [Timestamp('2016-08-19', tz='UTC'), Timestamp('2016-09-19', tz='UTC'), - Timestamp('2016-10-19', tz='UTC')], + Timestamp('2016-10-19', tz='UTC'), + Timestamp('2016-11-19', tz='UTC'), + Timestamp('2022-08-19', tz='UTC')], 'notice_date': [Timestamp('2016-01-26', tz='UTC'), Timestamp('2016-02-26', tz='UTC'), - Timestamp('2016-03-26', tz='UTC')], + Timestamp('2016-03-26', tz='UTC'), + Timestamp('2016-04-26', tz='UTC'), + Timestamp('2022-01-26', tz='UTC')], 'expiration_date': [Timestamp('2016-01-26', tz='UTC'), Timestamp('2016-02-26', tz='UTC'), - Timestamp('2016-03-26', tz='UTC')], + Timestamp('2016-03-26', tz='UTC'), + Timestamp('2016-04-26', tz='UTC'), + Timestamp('2022-01-26', tz='UTC')], 'auto_close_date': [Timestamp('2016-01-26', tz='UTC'), Timestamp('2016-02-26', tz='UTC'), - Timestamp('2016-03-26', tz='UTC')], - 'tick_size': [0.001] * 3, - 'multiplier': [1000.0] * 3, - 'exchange': ['CME'] * 3, + Timestamp('2016-03-26', tz='UTC'), + Timestamp('2016-04-26', tz='UTC'), + Timestamp('2022-01-26', tz='UTC')], + 'tick_size': [0.001] * 5, + 'multiplier': [1000.0] * 5, + 'exchange': ['CME'] * 5, }) def test_create_continuous_future(self): @@ -106,20 +118,20 @@ class ContinuousFuturesTestCase(WithCreateBarData, lambda: pd.Timestamp('2016-01-25', tz='UTC')) contract = bar_data.current(cf_primary, 'contract') - self.assertEqual(contract.symbol, 'FOF') + self.assertEqual(contract.symbol, 'FOF16') bar_data = self.create_bardata( lambda: pd.Timestamp('2016-01-26', tz='UTC')) contract = bar_data.current(cf_primary, 'contract') - self.assertEqual(contract.symbol, 'FOG', - 'Auto close at beginning of session so FOG is now ' + self.assertEqual(contract.symbol, 'FOG16', + 'Auto close at beginning of session so FOG16 is now ' 'the current contract.') bar_data = self.create_bardata( lambda: pd.Timestamp('2016-01-27', tz='UTC')) contract = bar_data.current(cf_primary, 'contract') - self.assertEqual(contract.symbol, 'FOG') + self.assertEqual(contract.symbol, 'FOG16') def test_current_contract_in_algo(self): code = dedent(""" @@ -145,28 +157,197 @@ def record_current_contract(algo, data): trading_calendar=self.trading_calendar, env=self.env) results = algo.run(self.data_portal) + result = results.iloc[0] - self.assertEqual(results.iloc[0].primary.symbol, - 'FOF', - 'Primary should be FOF on first session.') - self.assertEqual(results.iloc[0].secondary.symbol, - 'FOG', - 'Secondary should be FOG on first session.') + self.assertEqual(result.primary.symbol, + 'FOF16', + 'Primary should be FOF16 on first session.') + self.assertEqual(result.secondary.symbol, + 'FOG16', + 'Secondary should be FOG16 on first session.') + result = results.iloc[1] # Second day, primary should switch to FOG - self.assertEqual(results.iloc[1].primary.symbol, - 'FOG', - 'Primary should be FOG on second session, auto close ' - 'is at beginning of the session.') - self.assertEqual(results.iloc[1].secondary.symbol, - 'FOH', - 'Secondary should be FOH on second session, auto ' + self.assertEqual(result.primary.symbol, + 'FOG16', + 'Primary should be FOG16 on second session, auto ' + 'close is at beginning of the session.') + self.assertEqual(result.secondary.symbol, + 'FOH16', + 'Secondary should be FOH16 on second session, auto ' 'close is at beginning of the session.') + result = results.iloc[2] # Second day, primary should switch to FOG - self.assertEqual(results.iloc[2].primary.symbol, - 'FOG', - 'Primary should remain as FOG on third session.') - self.assertEqual(results.iloc[2].secondary.symbol, - 'FOH', - 'Secondary should remain as FOG on third session.') + self.assertEqual(result.primary.symbol, + 'FOG16', + 'Primary should remain as FOG16 on third session.') + self.assertEqual(result.secondary.symbol, + 'FOH16', + 'Secondary should remain as FOH16 on third session.') + + def test_current_chain_in_algo(self): + code = dedent(""" +from zipline.api import ( + record, + continuous_future, + schedule_function, + get_datetime, +) + +def initialize(algo): + algo.primary_cl = continuous_future('FO', 0, 'calendar') + algo.secondary_cl = continuous_future('FO', 1, 'calendar') + schedule_function(record_current_contract) + +def record_current_contract(algo, data): + record(datetime=get_datetime()) + primary_chain = data.current_chain(algo.primary_cl) + secondary_chain = data.current_chain(algo.secondary_cl) + record(primary_len=len(primary_chain)) + record(primary_first=primary_chain[0].symbol) + record(primary_last=primary_chain[-1].symbol) + record(secondary_len=len(secondary_chain)) + record(secondary_first=secondary_chain[0].symbol) + record(secondary_last=secondary_chain[-1].symbol) +""") + algo = TradingAlgorithm(script=code, + sim_params=self.sim_params, + trading_calendar=self.trading_calendar, + env=self.env) + results = algo.run(self.data_portal) + result = results.iloc[0] + + self.assertEqual(result.primary_len, + 4, + 'There should be only 4 contracts in the chain for ' + 'the primary, there are 5 contracts defined in the ' + 'fixture, but one has a start after the simulation ' + 'date.') + self.assertEqual(result.secondary_len, + 3, + 'There should be only 3 contracts in the chain for ' + 'the primary, there are 5 contracts defined in the ' + 'fixture, but one has a start after the simulation ' + 'date. And the first is not included because it is ' + 'the primary on that date.') + + self.assertEqual(result.primary_first, + 'FOF16', + 'Front of primary chain should be FOF16 on first ' + 'session.') + self.assertEqual(result.secondary_first, + 'FOG16', + 'Front of secondary chain should be FOG16 on first ' + 'session.') + + self.assertEqual(result.primary_last, + 'FOJ16', + 'End of primary chain should be FOJ16 on first ' + 'session.') + self.assertEqual(result.secondary_last, + 'FOJ16', + 'End of secondary chain should be FOJ16 on first ' + 'session.') + + # Second day, primary should switch to FOG + result = results.iloc[1] + + self.assertEqual(result.primary_len, + 3, + 'There should be only 3 contracts in the chain for ' + 'the primary, there are 5 contracts defined in the ' + 'fixture, but one has a start after the simulation ' + 'date. The first is not included because of roll.') + self.assertEqual(result.secondary_len, + 2, + 'There should be only 2 contracts in the chain for ' + 'the primary, there are 5 contracts defined in the ' + 'fixture, but one has a start after the simulation ' + 'date. The first is not included because of roll, ' + 'the second is the primary on that date.') + + self.assertEqual(result.primary_first, + 'FOG16', + 'Front of primary chain should be FOG16 on second ' + 'session.') + self.assertEqual(result.secondary_first, + 'FOH16', + 'Front of secondary chain should be FOH16 on second ' + 'session.') + + # These values remain FOJ16 because fixture data is not exhaustive + # enough to move the end of the chain. + self.assertEqual(result.primary_last, + 'FOJ16', + 'End of primary chain should be FOJ16 on second ' + 'session.') + self.assertEqual(result.secondary_last, + 'FOJ16', + 'End of secondary chain should be FOJ16 on second ' + 'session.') + + +class OrderedContractsTestCase(ZiplineTestCase): + + def test_active_chain(self): + contract_sids = array([1, 2, 3, 4], dtype=int64) + start_dates = pd.date_range('2015-01-01', periods=4, tz="UTC") + auto_close_dates = pd.date_range('2016-04-01', periods=4, tz="UTC") + + oc = OrderedContracts('FO', + contract_sids, + start_dates.astype('int64'), + auto_close_dates.astype('int64')) + + # Test sid 1 as days increment, as the sessions march forward + # a contract should be added per day, until all defined contracts + # are returned. + chain = oc.active_chain(1, pd.Timestamp('2014-12-31', tz='UTC').value) + self.assertEquals([], list(chain), + "On session before first start date, no contracts " + "in chain should be active.") + chain = oc.active_chain(1, pd.Timestamp('2015-01-01', tz='UTC').value) + self.assertEquals([1], list(chain), + "[1] should be the active chain on 01-01, since all " + "other start dates occur after 01-01.") + + chain = oc.active_chain(1, pd.Timestamp('2015-01-02', tz='UTC').value) + self.assertEquals([1, 2], list(chain), + "[1, 2] should be the active contracts on 01-02.") + + chain = oc.active_chain(1, pd.Timestamp('2015-01-03', tz='UTC').value) + self.assertEquals([1, 2, 3], list(chain), + "[1, 2, 3] should be the active contracts on 01-03.") + + chain = oc.active_chain(1, pd.Timestamp('2015-01-04', tz='UTC').value) + self.assertEquals(4, len(chain), + "[1, 2, 3, 4] should be the active contracts on " + "01-04, this is all defined contracts in the test " + "case.") + + chain = oc.active_chain(1, pd.Timestamp('2015-01-05', tz='UTC').value) + self.assertEquals(4, len(chain), + "[1, 2, 3, 4] should be the active contracts on " + "01-05. This tests the case where all start dates " + "are before the query date.") + + # Test querying each sid at a time when all should be alive. + chain = oc.active_chain(2, pd.Timestamp('2015-01-05', tz='UTC').value) + self.assertEquals([2, 3, 4], list(chain)) + + chain = oc.active_chain(3, pd.Timestamp('2015-01-05', tz='UTC').value) + self.assertEquals([3, 4], list(chain)) + + chain = oc.active_chain(4, pd.Timestamp('2015-01-05', tz='UTC').value) + self.assertEquals([4], list(chain)) + + # Test defined contract to check edge conditions. + chain = oc.active_chain(4, pd.Timestamp('2015-01-03', tz='UTC').value) + self.assertEquals([], list(chain), + "No contracts should be active, since 01-03 is " + "before 4's start date.") + + chain = oc.active_chain(4, pd.Timestamp('2015-01-04', tz='UTC').value) + self.assertEquals([4], list(chain), + "[4] should be active beginning at its start date.") diff --git a/zipline/_protocol.pyx b/zipline/_protocol.pyx index a9ee3ad6..c90a5224 100644 --- a/zipline/_protocol.pyx +++ b/zipline/_protocol.pyx @@ -431,6 +431,13 @@ cdef class BarData: return pd.DataFrame(data) + @check_parameters(('continuous_future',), + (ContinuousFuture,)) + def current_chain(self, continuous_future): + return self.data_portal.get_current_future_chain( + continuous_future, + self.simulation_dt_func()) + @check_parameters(('assets',), (Asset,)) def can_trade(self, assets): """ diff --git a/zipline/assets/continuous_futures.pyx b/zipline/assets/continuous_futures.pyx index 1ab2c054..2896403c 100644 --- a/zipline/assets/continuous_futures.pyx +++ b/zipline/assets/continuous_futures.pyx @@ -29,10 +29,9 @@ from cpython.object cimport ( ) from cpython cimport bool -import numpy as np +from numpy import empty from numpy cimport long_t, int64_t import warnings -cimport numpy as np from zipline.utils.calendars import get_calendar @@ -273,7 +272,7 @@ cdef class OrderedContracts(object): self.start_dates = start_dates self.auto_close_dates = auto_close_dates - cpdef long contract_before_auto_close(self, long_t dt_value): + cpdef long_t contract_before_auto_close(self, long_t dt_value): """ Get the contract with next upcoming auto close date. """ @@ -283,7 +282,7 @@ cdef class OrderedContracts(object): break return self.contract_sids[i] - cpdef long contract_at_offset(self, long_t sid, Py_ssize_t offset): + cpdef long_t 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. @@ -294,3 +293,23 @@ cdef class OrderedContracts(object): for i in range(self._size): if sid == sids[i]: return sids[i + offset] + + cpdef long_t[:] active_chain(self, long_t starting_sid, long_t dt_value): + cdef Py_ssize_t left, right, i, j + cdef long_t[:] sids, start_dates + left = 0 + right = self._size + sids = self.contract_sids + start_dates = self.start_dates + + for i in range(self._size): + if starting_sid == sids[i]: + left = i + break + + for j in range(i, self._size): + if start_dates[j] > dt_value: + right = j + break + + return sids[left:right] diff --git a/zipline/data/data_portal.py b/zipline/data/data_portal.py index ce543887..58fdc91b 100644 --- a/zipline/data/data_portal.py +++ b/zipline/data/data_portal.py @@ -1239,6 +1239,27 @@ class DataPortal(object): return ret + def get_current_future_chain(self, continuous_future, dt): + """ + Retrieves the future chain for the contract at the given `dt` according + the `continuous_future` specification. + + Returns: + future_chain : list[Future] + A list of active futures, where the first index is the current + contract specified by the continuous future definition, the second + is the next upcoming contract and so on. + """ + rf = self._roll_finders[continuous_future.roll_style] + session = self.trading_calendar.minute_to_session_label(dt) + contract_center = rf.get_contract_center( + continuous_future.root_symbol, session, + continuous_future.offset) + oc = self.asset_finder.get_ordered_contracts( + continuous_future.root_symbol) + chain = oc.active_chain(contract_center, session.value) + return self.asset_finder.retrieve_all(chain) + def _get_current_contract(self, continuous_future, dt): rf = self._roll_finders[continuous_future.roll_style] return self.asset_finder.retrieve_asset(