Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: Add vectorized lookup_symbol. #1627

Merged
merged 3 commits into from
Dec 28, 2016
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions tests/test_assets.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@
from zipline.testing import (
all_subindices,
empty_assets_db,
parameter_space,
tmp_assets_db,
)
from zipline.testing.predicates import assert_equal
Expand Down Expand Up @@ -1475,3 +1476,78 @@ def select_fields(r):
))

assert_equal(expected_data, actual_data)


class TestVectorizedSymbolLookup(WithAssetFinder, ZiplineTestCase):

@classmethod
def make_equity_info(cls):
T = partial(pd.Timestamp, tz='UTC')

def asset(sid, symbol, start_date, end_date):
return dict(
sid=sid,
symbol=symbol,
start_date=T(start_date),
end_date=T(end_date),
exchange='NYSE',
exchange_full='NYSE',
)

records = [
asset(1, 'A', '2014-01-02', '2014-01-31'),
asset(2, 'A', '2014-02-03', '2015-01-02'),
asset(3, 'B', '2014-01-02', '2014-01-15'),
asset(4, 'B', '2014-01-17', '2015-01-02'),
asset(5, 'C', '2001-01-02', '2015-01-02'),
asset(6, 'D', '2001-01-02', '2015-01-02'),
asset(7, 'FUZZY', '2001-01-02', '2015-01-02'),
]
return pd.DataFrame.from_records(records)

@parameter_space(
as_of=pd.to_datetime([
'2014-01-02',
'2014-01-15',
'2014-01-17',
'2015-01-02',
], utc=True),
symbols=[
[],
['A'], ['B'], ['C'], ['D'],
list('ABCD'),
list('ABCDDCBA'),
list('AABBAABBACABD'),
],
)
def test_lookup_symbols(self, as_of, symbols):
af = self.asset_finder
expected = [
af.lookup_symbol(symbol, as_of) for symbol in symbols
]
result = af.lookup_symbols(symbols, as_of)
assert_equal(result, expected)

def test_fuzzy(self):
af = self.asset_finder

# FUZZ.Y shouldn't resolve unless fuzzy=True.
syms = ['A', 'B', 'FUZZ.Y']
dt = pd.Timestamp('2014-01-15', tz='UTC')

with self.assertRaises(SymbolNotFound):
af.lookup_symbols(syms, pd.Timestamp('2014-01-15', tz='UTC'))

with self.assertRaises(SymbolNotFound):
af.lookup_symbols(
syms,
pd.Timestamp('2014-01-15', tz='UTC'),
fuzzy=False,
)

results = af.lookup_symbols(syms, dt, fuzzy=True)
assert_equal(results, af.retrieve_all([1, 3, 7]))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should you compare against a comprehension again?

assert_equal(
results,
[af.lookup_symbol(sym, dt, fuzzy=True) for sym in syms],
)
34 changes: 34 additions & 0 deletions zipline/assets/assets.py
Original file line number Diff line number Diff line change
Expand Up @@ -827,6 +827,40 @@ def lookup_symbol(self, symbol, as_of_date, fuzzy=False):
return self._lookup_symbol_fuzzy(symbol, as_of_date)
return self._lookup_symbol_strict(symbol, as_of_date)

def lookup_symbols(self, symbols, as_of_date, fuzzy=False):
"""
Lookup a list of equities by symbol.

Equivalent to::

[finder.lookup_symbol(s, as_of, fuzzy) for s in symbols]

but potentially faster because repeated lookups are memoized.

Parameters
----------
symbols : sequence[str]
Sequence of ticker symbols to resolve.
as_of_date : pd.Timestamp
Forwarded to ``lookup_symbol``.
fuzzy : bool, optional
Forwarded to ``lookup_symbol``.

Returns
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe add a note saying that it is the same as the list comprehension but better for performance.

-------
equities : list[Equity]
"""
memo = {}
out = []
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this just be a generator?

Copy link
Contributor Author

@ssanderson ssanderson Dec 28, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you asking if the internal implementation should be a generator, or if the top-level interface should be a generator? I'm indifferent toward the former. w/r/t the latter, I can't think of a case where you'd want to call this and not have strict semantics, and returning a generator creates issues for using this with numpy/pandas, so I'd rather have this be eager.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was talking about the top level impl, but this the numpy case makes sense.

append_output = out.append
for sym in symbols:
if sym in memo:
append_output(memo[sym])
else:
equity = memo[sym] = self.lookup_symbol(sym, as_of_date, fuzzy)
append_output(equity)
return out

def lookup_future_symbol(self, symbol):
"""Lookup a future contract by symbol.

Expand Down