From cc77a523220c980f3c741b073a0f3eea5ca8db51 Mon Sep 17 00:00:00 2001 From: Andrew Daniels Date: Wed, 24 Jun 2015 13:47:54 -0400 Subject: [PATCH] ENH: Adds future chain cache and future lookups to AssetFinder --- tests/test_assets.py | 122 +++++++++++++++++++++++++++++++++++++++ zipline/assets/assets.py | 111 +++++++++++++++++++++++++++++++++++ 2 files changed, 233 insertions(+) diff --git a/tests/test_assets.py b/tests/test_assets.py index 8e255fa0..1c8f3df9 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -574,3 +574,125 @@ class AssetFinderTestCase(TestCase): for warning in w: self.assertTrue(issubclass(warning.category, DeprecationWarning)) + + def test_lookup_future_in_chain(self): + metadata = { + 2: { + 'symbol': 'ADN15', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-06-15', tz='UTC') + }, + 1: { + 'symbol': 'ADV15', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-09-14', tz='UTC') + }, + 0: { + 'symbol': 'ADF16', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-12-14', tz='UTC') + }, + + } + + finder = AssetFinder(metadata=metadata) + dt = pd.Timestamp('2015-06-19', tz='UTC') + + # Check that the primary and secondary contracts are as expected + primary = finder.lookup_future_in_chain('AD', dt, 1) + secondary = finder.lookup_future_in_chain('AD', dt, 2) + self.assertEqual(primary.sid, 1) + self.assertEqual(secondary.sid, 0) + + # Check that we get None for an invalid contract num + self.assertIsNone(finder.lookup_future_in_chain('AD', dt, 0)) + self.assertIsNone(finder.lookup_future_in_chain('AD', dt, -10)) + self.assertIsNone(finder.lookup_future_in_chain('AD', dt, 10)) + self.assertIsNone(finder.lookup_future_in_chain('CL', dt, 1)) + + def test_lookup_future_chain(self): + metadata = { + 2: { + 'symbol': 'ADN15', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-06-15', tz='UTC') + }, + 1: { + 'symbol': 'ADV15', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-09-14', tz='UTC') + }, + 0: { + 'symbol': 'ADF16', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-12-14', tz='UTC') + }, + + } + + finder = AssetFinder(metadata=metadata) + dt = pd.Timestamp('2015-06-19', tz='UTC') + + # Check that we get the expected number of contract, in the + # right order + ad_contracts = finder.lookup_future_chain('AD', dt) + self.assertEqual(len(ad_contracts), 2) + self.assertEqual(ad_contracts[0].sid, 1) + self.assertEqual(ad_contracts[1].sid, 0) + + def test_lookup_future_by_expiration(self): + metadata = { + 2: { + 'symbol': 'ADN15', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-06-15', tz='UTC') + }, + 1: { + 'symbol': 'ADV15', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-09-14', tz='UTC') + }, + 0: { + 'symbol': 'ADF16', + 'root_symbol': 'AD', + 'asset_type': 'future', + 'expiration_date': pd.Timestamp('2015-12-14', tz='UTC') + }, + + } + + finder = AssetFinder(metadata=metadata) + dt = pd.Timestamp('2015-06-19', tz='UTC') + + # First-of-the-month timestamps + may_15 = pd.Timestamp('2015-05-01', tz='UTC') + june_15 = pd.Timestamp('2015-06-01', tz='UTC') + sept_15 = pd.Timestamp('2015-09-01', tz='UTC') + dec_15 = pd.Timestamp('2015-12-01', tz='UTC') + jan_16 = pd.Timestamp('2016-01-01', tz='UTC') + + # ADV15 is the next valid contract, so check that we get it + # for every ref_date before 9/14/15 + contract = finder.lookup_future_by_expiration('AD', dt, may_15) + self.assertEqual(contract.sid, 1) + + contract = finder.lookup_future_by_expiration('AD', dt, june_15) + self.assertEqual(contract.sid, 1) + + contract = finder.lookup_future_by_expiration('AD', dt, sept_15) + self.assertEqual(contract.sid, 1) + + # ADF16 has the next expiration date after 12/1/15 + contract = finder.lookup_future_by_expiration('AD', dt, dec_15) + self.assertEqual(contract.sid, 0) + + # No contracts exist after 12/14/2015, so we should get none + self.assertIsNone(finder.lookup_future_by_expiration('AD', dt, jan_16)) diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 43967f6f..f100a2ec 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -67,6 +67,7 @@ class AssetFinder(object): self.cache = {} self.sym_cache = {} + self.future_chains_cache = {} self.identifier_cache = {} self.fuzzy_match = {} @@ -230,6 +231,107 @@ class AssetFinder(object): else: self.fuzzy_match[(symbol, fuzzy, as_of_date)] = None + def _sort_future_chains(self): + """ Sort by increasing expiration date the list of contracts + for each root symbol in the future cache. + """ + exp_key = operator.attrgetter('expiration_date') + + for root_symbol in self.future_chains_cache: + self.future_chains_cache[root_symbol].sort(key=exp_key) + + def _valid_contracts(self, root_symbol, as_of_date): + """ Returns a list of the currently valid futures contracts + for a given root symbol, sorted by expiration date (the + contracts are sorted when the AssetFinder is built). + """ + try: + return [c for c in self.future_chains_cache[root_symbol] + if c.expiration_date > as_of_date] + except KeyError: + return None + + def lookup_future_chain(self, root_symbol, as_of_date): + """ Return the futures chain for a given root symbol. + + Parameters + ---------- + root_symbol : str + Root symbol of the desired future. + as_of_date : pd.Timestamp + Date at the time of the lookup. + + Returns + ------- + [Future] + """ + root_symbol.upper() + as_of_date = normalize_date(as_of_date) + return self._valid_contracts(root_symbol, as_of_date) + + def lookup_future_in_chain(self, root_symbol, as_of_date, contract_num=1): + """ Find a specific contract in the futures chain for a given + root symbol. + + Parameters + ---------- + root_symbol : str + Root symbol of the desired future. + as_of_date : pd.Timestamp + Date at the time of the lookup. + contract_num : int + 1 for the primary contract, 2 for the secondary, etc., + relative to as_of_date. + + Returns + ------- + Future + The (contract_num)th contract in the futures chain. If none + exits, returns None. + """ + root_symbol.upper() + as_of_date = normalize_date(as_of_date) + + valid_contracts = self._valid_contracts(root_symbol, as_of_date) + contract_index = contract_num - 1 + + if valid_contracts and contract_index >= 0: + try: + return valid_contracts[contract_index] + except IndexError: + pass + + return None + + def lookup_future_by_expiration(self, root_symbol, as_of_date, ref_date): + """ Find a specific contract in the futures chain by expiration + date. + + Parameters + ---------- + root_symbol : str + Root symbol of the desired future. + as_of_date : pd.Timestamp + Date at the time of the lookup. + ref_date : pd.Timestamp + Reference point for expiration dates. + + Returns + ------- + Future + The valid contract the has the closest expiration date + after ref_date. If none exists, returns None. + """ + root_symbol.upper() + as_of_date = normalize_date(as_of_date) + ref_date = normalize_date(ref_date) + + valid_contracts = self._valid_contracts(root_symbol, as_of_date) + + contracts_after_date = (c for c in valid_contracts + if c.expiration_date > ref_date) + return next(contracts_after_date, None) + def populate_cache(self): """ Populates the asset cache with all values in the assets @@ -239,6 +341,7 @@ class AssetFinder(object): # Wipe caches before repopulating self.cache = {} self.sym_cache = {} + self.future_chains_cache = {} self.identifier_cache = {} self.fuzzy_match = {} @@ -252,6 +355,14 @@ class AssetFinder(object): if asset.symbol is not '': self.sym_cache.setdefault(asset.symbol, []).append(asset) + if isinstance(asset, Future) and asset.root_symbol is not '': + self.future_chains_cache.setdefault(asset.root_symbol, + []).append(asset) + + # Pre-sort the future chains, we assume in future lookups + # that they're ordered correctly. + self._sort_future_chains() + def _spawn_asset(self, identifier, **kwargs): # Check if the sid is declared