diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 5443f0f6..52b78ea6 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -37,6 +37,7 @@ from zipline.errors import ( TradingControlViolation, AccountControlViolation, SymbolNotFound, + RootSymbolNotFound, ) from zipline.test_algorithms import ( access_account_in_init, @@ -335,6 +336,74 @@ class TestMiscellaneousAPI(TestCase): self.assertIsInstance(algo.sid(0), Equity) self.assertIsInstance(algo.sid(1), Equity) + def test_future_chain(self): + """ Tests the future_chain API function. + """ + + metadata = { + 0: { + 'symbol': 'CLG06', + 'root_symbol': 'CL', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2005-12-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-01-20', tz='UTC')}, + 1: { + 'root_symbol': 'CL', + 'symbol': 'CLK06', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2006-03-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-04-20', tz='UTC')}, + 2: { + 'symbol': 'CLQ06', + 'root_symbol': 'CL', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2006-06-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-07-20', tz='UTC')}, + 3: { + 'symbol': 'CLX06', + 'root_symbol': 'CL', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2006-02-01', tz='UTC'), + 'notice_date': pd.Timestamp('2006-09-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-10-20', tz='UTC')} + } + + algo = TradingAlgorithm(asset_metadata=metadata) + algo.datetime = pd.Timestamp('2006-12-01', tz='UTC') + + # Check that the fields of the FutureChain object are set correctly + cl = algo.future_chain('CL') + self.assertEqual(cl.root_symbol, 'CL') + self.assertEqual(cl.as_of_date, algo.datetime) + + # Check the fields are set correctly if an as_of_date is supplied + as_of_date = pd.Timestamp('1952-08-11', tz='UTC') + + cl = algo.future_chain('CL', as_of_date=as_of_date) + self.assertEqual(cl.root_symbol, 'CL') + self.assertEqual(cl.as_of_date, as_of_date) + + cl = algo.future_chain('CL', as_of_date='1952-08-11') + self.assertEqual(cl.root_symbol, 'CL') + self.assertEqual(cl.as_of_date, as_of_date) + + # Check that weird capitalization is corrected + cl = algo.future_chain('cL') + self.assertEqual(cl.root_symbol, 'CL') + + cl = algo.future_chain('cl') + self.assertEqual(cl.root_symbol, 'CL') + + # Check that invalid root symbols raise RootSymbolNotFound + with self.assertRaises(RootSymbolNotFound): + algo.future_chain('CLZ') + + with self.assertRaises(RootSymbolNotFound): + algo.future_chain('') + class TestTransformAlgorithm(TestCase): def setUp(self): diff --git a/tests/test_assets.py b/tests/test_assets.py index 1f014d51..5c2c420d 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -32,10 +32,12 @@ import pandas as pd from nose_parameterized import parameterized from zipline.assets import Asset, Equity, Future, AssetFinder +from zipline.assets.futures import FutureChain from zipline.errors import ( SymbolNotFound, MultipleSymbolsFound, SidAssignmentError, + RootSymbolNotFound, ) @@ -617,3 +619,198 @@ class AssetFinderTestCase(TestCase): pre_map = [asset201, asset2, asset200, asset1] post_map = finder.map_identifier_index_to_sids(pre_map, dt) self.assertListEqual([201, 2, 200, 1], post_map) + + +class TestFutureChain(TestCase): + metadata = { + 0: { + 'symbol': 'CLG06', + 'root_symbol': 'CL', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2005-12-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-01-20', tz='UTC')}, + 1: { + 'root_symbol': 'CL', + 'symbol': 'CLK06', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2006-03-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-04-20', tz='UTC')}, + 2: { + 'symbol': 'CLQ06', + 'root_symbol': 'CL', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2005-12-01', tz='UTC'), + 'notice_date': pd.Timestamp('2006-06-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-07-20', tz='UTC')}, + 3: { + 'symbol': 'CLX06', + 'root_symbol': 'CL', + 'asset_type': 'future', + 'start_date': pd.Timestamp('2006-02-01', tz='UTC'), + 'notice_date': pd.Timestamp('2006-09-20', tz='UTC'), + 'expiration_date': pd.Timestamp('2006-10-20', tz='UTC')} + } + + asset_finder = AssetFinder(metadata=metadata) + + def test_len(self): + """ Test the __len__ method of FutureChain. + """ + # None of the contracts have started yet. + cl = FutureChain(self.asset_finder, lambda: '2005-11-30', 'CL') + self.assertEqual(len(cl), 0) + + # Sids 0, 1, & 2 have started, 3 has not yet started. + cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL') + self.assertEqual(len(cl), 3) + + # Sid 0 is still valid the day before its notice date. + cl = FutureChain(self.asset_finder, lambda: '2005-12-19', 'CL') + self.assertEqual(len(cl), 3) + + # Sid 0 is now invalid, leaving only Sids 1 & 2 valid. + cl = FutureChain(self.asset_finder, lambda: '2005-12-20', 'CL') + self.assertEqual(len(cl), 2) + + # Sid 3 has started, so 1, 2, & 3 are now valid. + cl = FutureChain(self.asset_finder, lambda: '2006-02-01', 'CL') + self.assertEqual(len(cl), 3) + + # All contracts are no longer valid. + cl = FutureChain(self.asset_finder, lambda: '2006-09-20', 'CL') + self.assertEqual(len(cl), 0) + + def test_getitem(self): + """ Test the __getitem__ method of FutureChain. + """ + cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL') + self.assertEqual(cl[0], 0) + self.assertEqual(cl[1], 1) + self.assertEqual(cl[2], 2) + with self.assertRaises(IndexError): + cl[3] + + cl = FutureChain(self.asset_finder, lambda: '2005-12-19', 'CL') + self.assertEqual(cl[0], 0) + + cl = FutureChain(self.asset_finder, lambda: '2005-12-20', 'CL') + self.assertEqual(cl[0], 1) + + cl = FutureChain(self.asset_finder, lambda: '2006-02-01', 'CL') + self.assertEqual(cl[-1], 3) + + def test_root_symbols(self): + """ Test that different variations on root symbols are handled + as expected. + """ + # Make sure this successfully gets the chain for CL. + cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL') + self.assertEqual(cl.root_symbol, 'CL') + + # These root symbols don't exist, so RootSymbolNotFound should + # be raised immediately. + with self.assertRaises(RootSymbolNotFound): + FutureChain(self.asset_finder, lambda: '2005-12-01', 'CLZ') + + with self.assertRaises(RootSymbolNotFound): + FutureChain(self.asset_finder, lambda: '2005-12-01', '') + + def test_repr(self): + """ Test the __repr__ method of FutureChain. + """ + cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL') + cl_feb = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL', + as_of_date='2006-02-01') + + # The default chain should not include the as of date. + self.assertEqual(repr(cl), "FutureChain(root_symbol='CL')") + + # An explicit as of date should show up in the repr. + self.assertEqual( + repr(cl_feb), + ("FutureChain(root_symbol='CL', " + "as_of_date='2006-02-01 00:00:00+00:00')") + ) + + def test_as_of(self): + """ Test the as_of method of FutureChain. + """ + cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL') + + # Test that the as_of_date is set correctly to the future + feb = '2006-02-01' + cl_feb = cl.as_of(feb) + self.assertEqual( + cl_feb.as_of_date, + pd.Timestamp(feb, tz='UTC') + ) + + # Test that the as_of_date is set correctly to the past, with + # args of str, datetime.datetime, and pd.Timestamp. + feb_prev = '2005-02-01' + cl_feb_prev = cl.as_of(feb_prev) + self.assertEqual( + cl_feb_prev.as_of_date, + pd.Timestamp(feb_prev, tz='UTC') + ) + + feb_prev = datetime(year=2005, month=2, day=1) + cl_feb_prev = cl.as_of(feb_prev) + self.assertEqual( + cl_feb_prev.as_of_date, + pd.Timestamp(feb_prev, tz='UTC') + ) + + feb_prev = pd.Timestamp('2005-02-01') + cl_feb_prev = cl.as_of(feb_prev) + self.assertEqual( + cl_feb_prev.as_of_date, + pd.Timestamp(feb_prev, tz='UTC') + ) + + # The chain as of the current dt should always be the same as + # the defualt chain. Tests date as str, pd.Timestamp, and + # datetime.datetime. + self.assertEqual(cl[0], cl.as_of('2005-12-01')[0]) + self.assertEqual(cl[0], cl.as_of(pd.Timestamp('2005-12-01'))[0]) + self.assertEqual( + cl[0], + cl.as_of(datetime(year=2005, month=12, day=1))[0] + ) + + def test_offset(self): + """ Test the offset method of FutureChain. + """ + cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL') + + # Test that an offset forward sets as_of_date as expected + self.assertEqual( + cl.offset('3 days').as_of_date, + cl.as_of_date + pd.Timedelta(days=3) + ) + + # Test that an offset backward sets as_of_date as expected, with + # time delta given as str, datetime.timedelta, and pd.Timedelta. + self.assertEqual( + cl.offset('-1000 days').as_of_date, + cl.as_of_date + pd.Timedelta(days=-1000) + ) + self.assertEqual( + cl.offset(timedelta(days=-1000)).as_of_date, + cl.as_of_date + pd.Timedelta(days=-1000) + ) + self.assertEqual( + cl.offset(pd.Timedelta('-1000 days')).as_of_date, + cl.as_of_date + pd.Timedelta(days=-1000) + ) + + # An offset of zero should give the original chain. + self.assertEqual(cl[0], cl.offset(0)[0]) + self.assertEqual(cl[0], cl.offset("0 days")[0]) + + # A string that doesn't represent a time delta should raise a + # ValueError. + with self.assertRaises(ValueError): + cl.offset("blah") diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 609373bd..43aea515 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -65,6 +65,7 @@ from zipline.finance.slippage import ( transact_partial ) from zipline.assets import Asset, Future +from zipline.assets.futures import FutureChain from zipline.gens.composites import date_sorted_sources from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.sources import DataFrameSource, DataPanelSource @@ -682,6 +683,36 @@ class TradingAlgorithm(object): """ return self.asset_finder.retrieve_asset(a_sid) + @api_method + def future_chain(self, root_symbol, as_of_date=None): + """ Look up a future chain with the specified parameters. + + Parameters + ---------- + root_symbol : str + The root symbol of a future chain. + as_of_date : datetime.datetime or pandas.Timestamp or str, optional + Date at which the chain determination is rooted. I.e. the + existing contract whose notice date is first after this date is + the primary contract, etc. + + Returns + ------- + FutureChain + The future chain matching the specified parameters. + + Raises + ------ + RootSymbolNotFound + If a future chain could not be found for the given root symbol. + """ + return FutureChain( + asset_finder=self.asset_finder, + get_datetime=self.get_datetime, + root_symbol=root_symbol.upper(), + as_of_date=as_of_date + ) + def _calculate_order_value_amount(self, asset, value): """ Calculates how many shares/contracts to order based on the type of diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 6d5a702d..153e5f1d 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -211,7 +211,7 @@ class AssetFinder(object): self.fuzzy_match[(symbol, fuzzy, as_of_date)] = None def _sort_future_chains(self): - """ Sort by increasing expiration date the list of contracts + """ Sort by increasing notice date the list of contracts for each root symbol in the future cache. """ notice_key = operator.attrgetter('notice_date') @@ -228,8 +228,8 @@ class AssetFinder(object): Root symbol of the desired future. as_of_date : pd.Timestamp Date at which the chain determination is rooted. I.e. the - existing contract that expires first after (or on) this date is - the primary contract, etc. + existing contract whose notice date is first after this + date is the primary contract, etc. knowledge_date : pd.Timestamp Date for determining which contracts exist for inclusion in this chain. Contracts exist only if they have a start_date @@ -237,7 +237,15 @@ class AssetFinder(object): Returns ------- - [Future] + list + A list of Future objects, the chain for the given + parameters. + + Raises + ------ + RootSymbolNotFound + Raised when a future chain could not be found for the given + root symbol. """ try: return [c for c in self.future_chains_cache[root_symbol] diff --git a/zipline/assets/futures.py b/zipline/assets/futures.py new file mode 100644 index 00000000..a10bb192 --- /dev/null +++ b/zipline/assets/futures.py @@ -0,0 +1,169 @@ +from pandas import Timestamp, Timedelta +from pandas.tseries.tools import normalize_date + + +class FutureChain(object): + """ Allows users to look up future contracts. + + Parameters + ---------- + asset_finder : AssetFinder + An AssetFinder for future contract lookups, in particular the + AssetFinder of the TradingAlgorithm instance. + get_datetime : function + A function that returns the simulation datetime, in particular + the get_datetime method of the TradingAlgorithm instance. + root_symbol : str + The root symbol of a future chain. + as_of_date : pandas.Timestamp, optional + Date at which the chain determination is rooted. I.e. the + existing contract whose notice date is first after this date is + the primary contract, etc. If not provided, the current + simulation date is used as the as_of_date. + + Attributes + ---------- + root_symbol : str + The root symbol of the future chain. + as_of_date + The current as-of date of this future chain. + + Methods + ------- + as_of(dt) + offset(time_delta) + + Raises + ------ + RootSymbolNotFound + Raised when the FutureChain is initialized with a root symbol for which + a future chain could not be found. + """ + def __init__(self, asset_finder, get_datetime, root_symbol, + as_of_date=None): + self.root_symbol = root_symbol + + # Reference to the algo's AssetFinder for contract lookups + self._asset_finder = asset_finder + # Reference to the algo's get_datetime to know the current dt + self._algorithm_get_datetime = get_datetime + + # If an as_of_date is provided, self._as_of_date uses that + # value, otherwise None. This attribute backs the as_of_date property. + if as_of_date: + self._as_of_date = normalize_date(Timestamp(as_of_date, tz='UTC')) + else: + self._as_of_date = None + + # Attribute to cache the most up-to-date chain, and the dt when it was + # last updated. + self._current_chain = [] + self._last_updated = None + + # Get the initial chain, since self._last_updated is None. + self._maybe_update_current_chain() + + def __repr__(self): + # NOTE: The string returned cannot be used to instantiate this + # exact FutureChain, since we don't want to display the asset + # finder and get_datetime function to the user. + if self._as_of_date: + return "FutureChain(root_symbol='%s', as_of_date='%s')" % ( + self.root_symbol, self.as_of_date) + else: + return "FutureChain(root_symbol='%s')" % self.root_symbol + + def _get_datetime(self): + """ + Returns the normalized simulation datetime. + + Returns + ------- + pandas.Timestamp + The normalized datetime of FutureChain's TradingAlgorithm. + """ + return normalize_date( + Timestamp(self._algorithm_get_datetime(), tz='UTC') + ) + + @property + def as_of_date(self): + """ + The current as-of date of this future chain. + + Returns + ------- + pandas.Timestamp + The user-provided as_of_date if given, otherwise the + current datetime of the simulation. + """ + if self._as_of_date is not None: + return self._as_of_date + else: + return self._get_datetime() + + def _maybe_update_current_chain(self): + """ Updates the current chain if it's out of date, then returns + it. + + Returns + ------- + list + The up-to-date current chain, a list of Future objects. + """ + dt = self._get_datetime() + + if (self._last_updated is None) or (self._last_updated != dt): + self._current_chain = self._asset_finder.lookup_future_chain( + self.root_symbol, + self.as_of_date, + dt + ) + self._last_updated = dt + + return self._current_chain + + def __getitem__(self, key): + return self._maybe_update_current_chain()[key] + + def __len__(self): + return len(self._maybe_update_current_chain()) + + def __iter__(self): + return iter(self._maybe_update_current_chain()) + + def as_of(self, dt): + """ Get the future chain for this root symbol as of a specific date. + + Parameters + ---------- + dt : datetime.datetime or pandas.Timestamp or str, optional + The as_of_date for the new chain. + + Returns + ------- + FutureChain + + """ + return FutureChain( + asset_finder=self._asset_finder, + get_datetime=self._algorithm_get_datetime, + root_symbol=self.root_symbol, + as_of_date=dt + ) + + def offset(self, time_delta): + """ Get the future chain for this root symbol with a given + offset from the current as_of_date. + + Parameters + ---------- + time_delta : datetime.timedelta or pandas.Timedelta or str + The offset from the current as_of_date for the new chain. + + Returns + ------- + FutureChain + + """ + return self.as_of(self.as_of_date + Timedelta(time_delta))