Merge pull request #940 from quantopian/generalize-databazaar-stuff

MAINT: Move code to shared modules
This commit is contained in:
Joe Jevnik
2016-01-08 17:01:39 -05:00
7 changed files with 278 additions and 200 deletions
+17 -50
View File
@@ -27,57 +27,15 @@ from zipline.pipeline.loaders.blaze import (
TS_FIELD_NAME,
)
from zipline.utils.numpy_utils import make_datetime64D, np_NaT
from zipline.utils.tradingcalendar import trading_days
from zipline.utils.test_utils import (
make_simple_equity_info,
powerset,
tmp_asset_finder,
gen_calendars,
to_series,
num_days_in_range,
)
def _to_series(knowledge_dates, earning_dates):
"""
Helper for converting a dict of strings to a Series of datetimes.
This is just for making the test cases more readable.
"""
return pd.Series(
index=pd.to_datetime(knowledge_dates),
data=pd.to_datetime(earning_dates),
)
def num_days_in_range(dates, start, end):
"""
Return the number of days in `dates` between start and end, inclusive.
"""
start_idx, stop_idx = dates.slice_locs(start, end)
return stop_idx - start_idx
def gen_calendars():
"""
Generate calendars to use as inputs to test_compute_latest.
"""
start, stop = '2014-01-01', '2014-01-31'
all_dates = pd.date_range(start, stop, tz='utc')
# These dates are the points where announcements or knowledge dates happen.
# Test every combination of them being absent.
critical_dates = pd.to_datetime([
'2014-01-05',
'2014-01-10',
'2014-01-15',
'2014-01-20',
])
for to_drop in map(list, powerset(critical_dates)):
# Have to yield tuples.
yield (all_dates.drop(to_drop),)
# Also test with the trading calendar.
yield (trading_days[trading_days.slice_indexer(start, stop)],)
class EarningsCalendarLoaderTestCase(TestCase):
"""
Tests for loading the earnings announcement data.
@@ -99,22 +57,22 @@ class EarningsCalendarLoaderTestCase(TestCase):
cls.earnings_dates = {
# K1--K2--E1--E2.
A: _to_series(
A: to_series(
knowledge_dates=['2014-01-05', '2014-01-10'],
earning_dates=['2014-01-15', '2014-01-20'],
),
# K1--K2--E2--E1.
B: _to_series(
B: to_series(
knowledge_dates=['2014-01-05', '2014-01-10'],
earning_dates=['2014-01-20', '2014-01-15']
),
# K1--E1--K2--E2.
C: _to_series(
C: to_series(
knowledge_dates=['2014-01-05', '2014-01-15'],
earning_dates=['2014-01-10', '2014-01-20']
),
# K1 == K2.
D: _to_series(
D: to_series(
knowledge_dates=['2014-01-05'] * 2,
earning_dates=['2014-01-10', '2014-01-15'],
),
@@ -294,7 +252,16 @@ class EarningsCalendarLoaderTestCase(TestCase):
index=announcement_dates.index,
)
@parameterized.expand(gen_calendars())
@parameterized.expand(gen_calendars(
'2014-01-01',
'2014-01-31',
critical_dates=pd.to_datetime([
'2014-01-05',
'2014-01-10',
'2014-01-15',
'2014-01-20',
]),
))
def test_compute_earnings(self, dates):
(
+84 -1
View File
@@ -154,7 +154,7 @@ from toolz import (
memoize,
)
import toolz.curried.operator as op
from six import with_metaclass, PY2, itervalues
from six import with_metaclass, PY2, itervalues, iteritems
from zipline.pipeline.data.dataset import DataSet, Column
@@ -903,3 +903,86 @@ class BlazeLoader(dict):
)
global_loader = BlazeLoader.global_instance()
def bind_expression_to_resources(expr, resources):
"""
Bind a Blaze expression to resources.
Parameters
----------
expr : bz.Expr
The expression to which we want to bind resources.
resources : dict[bz.Symbol -> any]
Mapping from the atomic terms of ``expr`` to actual data resources.
Returns
-------
bound_expr : bz.Expr
``expr`` with bound resources.
"""
# bind the resources into the expression
if resources is None:
resources = {}
# _subs stands for substitute. It's not actually private, blaze just
# prefixes symbol-manipulation methods with underscores to prevent
# collisions with data column names.
return expr._subs({
k: bz.Data(v, dshape=k.dshape) for k, v in iteritems(resources)
})
def ffill_query_in_range(expr,
lower,
upper,
odo_kwargs=None,
ts_field=TS_FIELD_NAME,
sid_field=SID_FIELD_NAME):
"""Query a blaze expression in a given time range properly forward filling
from values that fall before the lower date.
Parameters
----------
expr : Expr
Bound blaze expression.
lower : datetime
The lower date to query for.
upper : datetime
The upper date to query for.
odo_kwargs : dict, optional
The extra keyword arguments to pass to ``odo``.
ts_field : str, optional
The name of the timestamp field in the given blaze expression.
sid_field : str, optional
The name of the sid field in the given blaze expression.
Returns
-------
raw : pd.DataFrame
A strict dataframe for the data in the given date range. This may
start before the requested start date if a value is needed to ffill.
"""
odo_kwargs = odo_kwargs or {}
filtered = expr[expr[ts_field] <= lower]
computed_lower = odo(
bz.by(
filtered[sid_field],
timestamp=filtered[ts_field].max(),
).timestamp.min(),
pd.Timestamp,
**odo_kwargs
)
if pd.isnull(computed_lower):
# If there is no lower date, just query for data in the date
# range. It must all be null anyways.
computed_lower = lower
return odo(
expr[
(expr[ts_field] >= computed_lower) &
(expr[ts_field] <= upper)
],
pd.DataFrame,
**odo_kwargs
)
+14 -59
View File
@@ -1,11 +1,13 @@
import blaze as bz
from datashape import istabular
from odo import odo
import pandas as pd
from six import iteritems
from toolz import valmap
from .core import TS_FIELD_NAME, SID_FIELD_NAME
from .core import (
TS_FIELD_NAME,
SID_FIELD_NAME,
bind_expression_to_resources,
ffill_query_in_range,
)
from zipline.pipeline.data import EarningsCalendar
from zipline.pipeline.loaders.base import PipelineLoader
from zipline.pipeline.loaders.earnings import EarningsCalendarLoader
@@ -14,34 +16,6 @@ from zipline.pipeline.loaders.earnings import EarningsCalendarLoader
ANNOUNCEMENT_FIELD_NAME = 'announcement_date'
def bind_expression_to_resources(expr, resources):
"""
Bind a Blaze expression to resources.
Parameters
----------
expr : bz.Expr
The expression to which we want to bind resources.
resources : dict[bz.Symbol -> any]
Mapping from the atomic terms of ``expr`` to actual data resources.
Returns
-------
bound_expr : bz.Expr
``expr`` with bound resources.
"""
# bind the resources into the expression
if resources is None:
resources = {}
# _subs stands for substitute. It's not actually private, blaze just
# prefixes symbol-manipulation methods with underscores to prevent
# collisions with data column names.
return expr._subs({
k: bz.Data(v, dshape=k.dshape) for k, v in iteritems(resources)
})
class BlazeEarningsCalendarLoader(PipelineLoader):
"""A pipeline loader for the ``EarningsCalendar`` dataset that loads
data from a blaze expression.
@@ -61,8 +35,8 @@ class BlazeEarningsCalendarLoader(PipelineLoader):
Dim * {{
{SID_FIELD_NAME}: int64,
{TS_FIELD_NAME}: datetime64,
{ANNOUNCEMENT_FIELD_NAME}: datetime64,
{TS_FIELD_NAME}: datetime,
{ANNOUNCEMENT_FIELD_NAME}: ?datetime,
}}
Where each row of the table is a record including the sid to identify the
@@ -87,7 +61,6 @@ class BlazeEarningsCalendarLoader(PipelineLoader):
def __init__(self,
expr,
resources=None,
compute_kwargs=None,
odo_kwargs=None,
dataset=EarningsCalendar):
dshape = expr.dshape
@@ -106,33 +79,15 @@ class BlazeEarningsCalendarLoader(PipelineLoader):
self._dataset = dataset
def load_adjusted_array(self, columns, dates, assets, mask):
expr = self._expr
filtered = expr[expr[TS_FIELD_NAME] <= dates[0]]
lower = odo(
bz.by(
filtered[SID_FIELD_NAME],
timestamp=filtered[TS_FIELD_NAME].max(),
).timestamp.min(),
pd.Timestamp,
**self._odo_kwargs
raw = ffill_query_in_range(
self._expr,
dates[0],
dates[-1],
self._odo_kwargs,
)
if pd.isnull(lower):
# If there is no lower date, just query for data in the date
# range. It must all be null anyways.
lower = dates[0]
raw = odo(
expr[
(expr[TS_FIELD_NAME] >= lower) &
(expr[TS_FIELD_NAME] <= dates[-1])
],
pd.DataFrame,
**self._odo_kwargs
)
sids = raw.loc[:, SID_FIELD_NAME]
raw.drop(
sids[~(sids.isin(assets) | sids.notnull())].index,
sids[~sids.isin(assets)].index,
inplace=True
)
+3 -87
View File
@@ -3,16 +3,14 @@ Reference implementation for EarningsCalendar loaders.
"""
from itertools import repeat
from numpy import full_like, full
import pandas as pd
from six import iteritems
from six.moves import zip
from toolz import merge
from .base import PipelineLoader
from .frame import DataFrameLoader
from .utils import next_date_frame, previous_date_frame
from ..data.earnings import EarningsCalendar
from zipline.utils.numpy_utils import np_NaT
from zipline.utils.memoize import lazyval
@@ -83,7 +81,7 @@ class EarningsCalendarLoader(PipelineLoader):
def next_announcement_loader(self):
return DataFrameLoader(
self.dataset.next_announcement,
next_earnings_date_frame(
next_date_frame(
self.all_dates,
self.announcement_dates,
),
@@ -94,7 +92,7 @@ class EarningsCalendarLoader(PipelineLoader):
def previous_announcement_loader(self):
return DataFrameLoader(
self.dataset.previous_announcement,
previous_earnings_date_frame(
previous_date_frame(
self.all_dates,
self.announcement_dates,
),
@@ -108,85 +106,3 @@ class EarningsCalendarLoader(PipelineLoader):
)
for column in columns
)
def next_earnings_date_frame(dates, announcement_dates):
"""
Make a DataFrame representing simulated next earnings dates.
Parameters
----------
dates : pd.DatetimeIndex.
The index of the returned DataFrame.
announcement_dates : dict[int -> pd.Series]
Dict mapping sids to an index of dates on which earnings were announced
for that sid.
Returns
-------
next_earnings: pd.DataFrame
A DataFrame representing, for each (label, date) pair, the first entry
in `earnings_calendars[label]` on or after `date`. Entries falling
after the last date in a calendar will have `np_NaT` as the result in
the output.
See Also
--------
previous_earnings_date_frame
"""
cols = {equity: full_like(dates, np_NaT) for equity in announcement_dates}
raw_dates = dates.values
for equity, earnings_dates in iteritems(announcement_dates):
data = cols[equity]
if not earnings_dates.index.is_monotonic_increasing:
earnings_dates = earnings_dates.sort_index()
# Iterate over the raw Series values, since we're comparing against
# numpy arrays anyway.
iterkv = zip(earnings_dates.index.values, earnings_dates.values)
for timestamp, announce_date in iterkv:
date_mask = (timestamp <= raw_dates) & (raw_dates <= announce_date)
value_mask = (announce_date <= data) | (data == np_NaT)
data[date_mask & value_mask] = announce_date
return pd.DataFrame(index=dates, data=cols)
def previous_earnings_date_frame(dates, announcement_dates):
"""
Make a DataFrame representing simulated next earnings dates.
Parameters
----------
dates : DatetimeIndex.
The index of the returned DataFrame.
announcement_dates : dict[int -> DatetimeIndex]
Dict mapping sids to an index of dates on which earnings were announced
for that sid.
Returns
-------
prev_earnings: pd.DataFrame
A DataFrame representing, for (label, date) pair, the first entry in
`announcement_dates[label]` strictly before `date`. Entries falling
before the first date in a calendar will have `NaT` as the result in
the output.
See Also
--------
next_earnings_date_frame
"""
sids = list(announcement_dates)
out = full((len(dates), len(sids)), np_NaT, dtype='datetime64[ns]')
dn = dates[-1].asm8
for col_idx, sid in enumerate(sids):
# announcement_dates[sid] is Series mapping knowledge_date to actual
# announcement date. We don't care about the knowledge date for
# computing previous earnings.
values = announcement_dates[sid].values
values = values[values <= dn]
out[dates.searchsorted(values), col_idx] = values
frame = pd.DataFrame(out, index=dates, columns=sids)
frame.ffill(inplace=True)
return frame
+95
View File
@@ -0,0 +1,95 @@
import numpy as np
import pandas as pd
from six import iteritems
from six.moves import zip
from zipline.utils.numpy_utils import np_NaT
def next_date_frame(dates, events_by_sid):
"""
Make a DataFrame representing the simulated next known date for an event.
Parameters
----------
dates : pd.DatetimeIndex.
The index of the returned DataFrame.
events_by_sid : dict[int -> pd.Series]
Dict mapping sids to a series of dates. Each k:v pair of the series
represents the date we learned of the event mapping to the date the
event will occur.
Returns
-------
next_events: pd.DataFrame
A DataFrame where each column is a security from `events_by_sid` where
the values are the dates of the next known event with the knowledge we
had on the date of the index. Entries falling after the last date will
have `NaT` as the result in the output.
See Also
--------
previous_date_frame
"""
cols = {
equity: np.full_like(dates, np_NaT) for equity in events_by_sid
}
raw_dates = dates.values
for equity, event_dates in iteritems(events_by_sid):
data = cols[equity]
if not event_dates.index.is_monotonic_increasing:
event_dates = event_dates.sort_index()
# Iterate over the raw Series values, since we're comparing against
# numpy arrays anyway.
iterkv = zip(event_dates.index.values, event_dates.values)
for knowledge_date, event_date in iterkv:
date_mask = (
(knowledge_date <= raw_dates) &
(raw_dates <= event_date)
)
value_mask = (event_date <= data) | (data == np_NaT)
data[date_mask & value_mask] = event_date
return pd.DataFrame(index=dates, data=cols)
def previous_date_frame(date_index, events_by_sid):
"""
Make a DataFrame representing simulated next earnings date_index.
Parameters
----------
date_index : DatetimeIndex.
The index of the returned DataFrame.
events_by_sid : dict[int -> DatetimeIndex]
Dict mapping sids to a series of dates. Each k:v pair of the series
represents the date we learned of the event mapping to the date the
event will occur.
Returns
-------
previous_events: pd.DataFrame
A DataFrame where each column is a security from `events_by_sid` where
the values are the dates of the previous event that occured on the date
of the index. Entries falling before the first date will have `NaT` as
the result in the output.
See Also
--------
next_date_frame
"""
sids = list(events_by_sid)
out = np.full((len(date_index), len(sids)), np_NaT, dtype='datetime64[ns]')
dn = date_index[-1].asm8
for col_idx, sid in enumerate(sids):
# events_by_sid[sid] is Series mapping knowledge_date to actual
# event_date. We don't care about the knowledge date for
# computing previous earnings.
values = events_by_sid[sid].values
values = values[values <= dn]
out[date_index.searchsorted(values), col_idx] = values
frame = pd.DataFrame(out, index=date_index, columns=sids)
frame.ffill(inplace=True)
return frame
+31 -3
View File
@@ -7,9 +7,8 @@ from weakref import WeakKeyDictionary
class lazyval(object):
"""
Decorator that marks that an attribute should not be computed until
needed, and that the value should be memoized.
"""Decorator that marks that an attribute of an instance should not be
computed until needed, and that the value should be memoized.
Example
-------
@@ -57,6 +56,35 @@ class lazyval(object):
del self._cache[instance]
class classlazyval(lazyval):
""" Decorator that marks that an attribute of a class should not be
computed until needed, and that the value should be memoized.
Example
-------
>>> from zipline.utils.memoize import classlazyval
>>> class C(object):
... count = 0
... @classlazyval
... def val(cls):
... cls.count += 1
... return "val"
...
>>> C.count
0
>>> C.val, C.count
('val', 1)
>>> C.val, C.count
('val', 1)
"""
# We don't reassign the name on the class to implement the caching because
# then we would need to use a metaclass to track the name of the
# descriptor.
def __get__(self, instance, owner):
return super(classlazyval, self).__get__(owner, owner)
def remember_last(f):
"""
Decorator that remembers the last computed value of a function and doesn't
+34
View File
@@ -26,6 +26,7 @@ from zipline.assets.asset_writer import AssetDBWriterFromDataFrame
from zipline.assets.futures import CME_CODE_TO_MONTH
from zipline.finance.order import ORDER_STATUS
from zipline.utils import security_list
from zipline.utils.tradingcalendar import trading_days
EPOCH = pd.Timestamp(0, tz='UTC')
@@ -630,3 +631,36 @@ def powerset(values):
Return the power set (i.e., the set of all subsets) of entries in `values`.
"""
return concat(combinations(values, i) for i in range(len(values) + 1))
def to_series(knowledge_dates, earning_dates):
"""
Helper for converting a dict of strings to a Series of datetimes.
This is just for making the test cases more readable.
"""
return pd.Series(
index=pd.to_datetime(knowledge_dates),
data=pd.to_datetime(earning_dates),
)
def num_days_in_range(dates, start, end):
"""
Return the number of days in `dates` between start and end, inclusive.
"""
start_idx, stop_idx = dates.slice_locs(start, end)
return stop_idx - start_idx
def gen_calendars(start, stop, critical_dates):
"""
Generate calendars to use as inputs.
"""
all_dates = pd.date_range(start, stop, tz='utc')
for to_drop in map(list, powerset(critical_dates)):
# Have to yield tuples.
yield (all_dates.drop(to_drop),)
# Also test with the trading calendar.
yield (trading_days[trading_days.slice_indexer(start, stop)],)