Merge pull request #1313 from nathanwolfe/master

BUG: Add support for Panel data in accordance with documentation
This commit is contained in:
Joe Jevnik
2016-07-29 20:11:56 -04:00
committed by GitHub
6 changed files with 201 additions and 62 deletions
+85 -1
View File
@@ -33,7 +33,10 @@ import numpy as np
import pandas as pd
import pytz
from zipline import TradingAlgorithm
from zipline import (
run_algorithm,
TradingAlgorithm,
)
from zipline.api import FixedSlippage
from zipline.assets import Equity, Future
from zipline.assets.synthetic import (
@@ -161,6 +164,7 @@ from zipline.test_algorithms import (
no_handle_data,
)
from zipline.utils.api_support import ZiplineAPI, set_algo_instance
from zipline.utils.calendars import get_calendar
from zipline.utils.context_tricks import CallbackManager
from zipline.utils.control_flow import nullctx
import zipline.utils.events
@@ -4102,3 +4106,83 @@ class AlgoInputValidationTestCase(ZiplineTestCase):
script=script,
**{method: lambda *args, **kwargs: None}
)
class TestPanelData(ZiplineTestCase):
@parameterized.expand([
('daily',
pd.Timestamp('2015-12-23', tz='UTC'),
pd.Timestamp('2016-01-05', tz='UTC'),),
('minute',
pd.Timestamp('2015-12-23', tz='UTC'),
pd.Timestamp('2015-12-24', tz='UTC'),),
])
def test_panel_data(self, data_frequency, start_dt, end_dt):
trading_calendar = get_calendar('NYSE')
if data_frequency == 'daily':
history_freq = '1d'
create_df_for_asset = create_daily_df_for_asset
dt_transform = trading_calendar.minute_to_session_label
elif data_frequency == 'minute':
history_freq = '1m'
create_df_for_asset = create_minute_df_for_asset
def dt_transform(dt):
return dt
sids = range(1, 3)
dfs = {}
for sid in sids:
dfs[sid] = create_df_for_asset(trading_calendar,
start_dt, end_dt, interval=sid)
dfs[sid]['prev_close'] = dfs[sid]['close'].shift(1)
panel = pd.Panel(dfs)
price_record = pd.Panel(items=sids,
major_axis=panel.major_axis,
minor_axis=['current', 'previous'])
def initialize(algo):
algo.first_bar = True
algo.equities = []
for sid in sids:
algo.equities.append(algo.sid(sid))
def handle_data(algo, data):
price_record.loc[:, dt_transform(algo.get_datetime()),
'current'] = (
data.current(algo.equities, 'price')
)
if algo.first_bar:
algo.first_bar = False
else:
price_record.loc[:, dt_transform(algo.get_datetime()),
'previous'] = (
data.history(algo.equities, 'price',
2, history_freq).iloc[0]
)
def check_panels():
np.testing.assert_array_equal(
price_record.values.astype('float64'),
panel.loc[:, :, ['close',
'prev_close']].values.astype('float64')
)
trading_algo = TradingAlgorithm(initialize=initialize,
handle_data=handle_data)
trading_algo.run(data=panel)
check_panels()
price_record.loc[:] = np.nan
run_algorithm(
start=start_dt,
end=end_dt,
capital_base=1,
initialize=initialize,
handle_data=handle_data,
data_frequency=data_frequency,
data=panel
)
check_panels()
@@ -18,31 +18,29 @@ from itertools import permutations, product
import numpy as np
import pandas as pd
from zipline.data.us_equity_pricing import PanelDailyBarReader
from zipline.data.us_equity_pricing import PanelBarReader
from zipline.testing import ExplodingObject
from zipline.testing.fixtures import (
WithAssetFinder,
WithNYSETradingDays,
ZiplineTestCase,
)
from zipline.utils.calendars import get_calendar
class TestPanelDailyBarReader(WithAssetFinder,
WithNYSETradingDays,
ZiplineTestCase):
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
END_DATE = pd.Timestamp('2006-02-01', tz='utc')
class WithPanelBarReader(WithAssetFinder):
@classmethod
def init_class_fixtures(cls):
super(TestPanelDailyBarReader, cls).init_class_fixtures()
super(WithPanelBarReader, cls).init_class_fixtures()
finder = cls.asset_finder
days = cls.trading_days
trading_calendar = get_calendar('NYSE')
items = finder.retrieve_all(finder.sids)
major_axis = days
major_axis = (
trading_calendar.sessions_in_range if cls.FREQUENCY == 'daily'
else trading_calendar.minutes_for_sessions_in_range
)(cls.START_DATE, cls.END_DATE)
minor_axis = ['open', 'high', 'low', 'close', 'volume']
shape = tuple(map(len, [items, major_axis, minor_axis]))
@@ -55,7 +53,7 @@ class TestPanelDailyBarReader(WithAssetFinder,
minor_axis=minor_axis,
)
cls.reader = PanelDailyBarReader(days, cls.panel)
cls.reader = PanelBarReader(trading_calendar, cls.panel, cls.FREQUENCY)
def test_spot_price(self):
panel = self.panel
@@ -83,7 +81,7 @@ class TestPanelDailyBarReader(WithAssetFinder,
for axis_order in permutations((0, 1, 2)):
transposed = panel.transpose(*axis_order)
with self.assertRaises(ValueError) as e:
PanelDailyBarReader(unused, transposed)
PanelBarReader(unused, transposed, 'daily')
expected = (
"Duplicate entries in Panel.{name}: ['a', 'b'].".format(
@@ -95,6 +93,28 @@ class TestPanelDailyBarReader(WithAssetFinder,
def test_sessions(self):
sessions = self.reader.sessions
self.assertEqual(21, len(sessions))
self.assertEqual(self.NUM_SESSIONS, len(sessions))
self.assertEqual(self.START_DATE, sessions[0])
self.assertEqual(self.END_DATE, sessions[-1])
class TestPanelDailyBarReader(WithPanelBarReader,
ZiplineTestCase):
FREQUENCY = 'daily'
START_DATE = pd.Timestamp('2006-01-03', tz='utc')
END_DATE = pd.Timestamp('2006-02-01', tz='utc')
NUM_SESSIONS = 21
class TestPanelMinuteBarReader(WithPanelBarReader,
ZiplineTestCase):
FREQUENCY = 'minute'
START_DATE = pd.Timestamp('2015-12-23', tz='utc')
END_DATE = pd.Timestamp('2015-12-24', tz='utc')
NUM_SESSIONS = 2
+30 -7
View File
@@ -37,7 +37,7 @@ from six import (
from zipline._protocol import handle_non_market_minutes
from zipline.assets.synthetic import make_simple_equity_info
from zipline.data.data_portal import DataPortal
from zipline.data.us_equity_pricing import PanelDailyBarReader
from zipline.data.us_equity_pricing import PanelBarReader
from zipline.errors import (
AttachPipelineAfterInitialize,
HistoryInInitialize,
@@ -611,14 +611,30 @@ class TradingAlgorithm(object):
data = data.swapaxes(0, 2)
if isinstance(data, pd.Panel):
# Guard against tz-naive index.
if data.major_axis.tz is None:
data.major_axis = data.major_axis.tz_localize('UTC')
# For compatibility with existing examples allow start/end
# to be inferred.
if overwrite_sim_params:
self.sim_params = self.sim_params.create_new(
data.major_axis[0],
data.major_axis[-1]
self.trading_calendar.minute_to_session_label(
data.major_axis[0]
),
self.trading_calendar.minute_to_session_label(
data.major_axis[-1]
),
)
# Assume data is daily if timestamp times are
# standardized, otherwise assume minute bars.
times = data.major_axis.time
if np.all(times == times[0]):
self.sim_params.data_frequency = 'daily'
else:
self.sim_params.data_frequency = 'minute'
copy_panel = data.rename(
# These were the old names for the close/open columns. We
# need to make a copy anyway, so swap these for backwards
@@ -634,15 +650,22 @@ class TradingAlgorithm(object):
copy_panel.items
)
)
equity_daily_reader = PanelDailyBarReader(
self.trading_calendar.all_sessions,
if self.sim_params.data_frequency == 'daily':
equity_reader_arg = 'equity_daily_reader'
elif self.sim_params.data_frequency == 'minute':
equity_reader_arg = 'equity_minute_reader'
equity_reader = PanelBarReader(
self.trading_calendar,
copy_panel,
self.sim_params.data_frequency,
)
self.data_portal = DataPortal(
self.asset_finder,
self.trading_calendar,
first_trading_day=equity_daily_reader.first_trading_day,
equity_daily_reader=equity_daily_reader,
first_trading_day=equity_reader.first_trading_day,
**{equity_reader_arg: equity_reader}
)
# Force a reset of the performance tracker, in case
-2
View File
@@ -156,8 +156,6 @@ class DataPortal(object):
self._equity_minute_reader,
self._adjustment_reader
)
self.MINUTE_PRICE_ADJUSTMENT_FACTOR = \
self._equity_minute_reader._ohlc_inverse
self._first_trading_day = first_trading_day
+40 -34
View File
@@ -35,15 +35,14 @@ from numpy import (
issubdtype,
nan,
uint32,
zeros,
)
from pandas import (
DataFrame,
read_csv,
Timestamp,
NaT,
isnull,
DatetimeIndex)
DatetimeIndex
)
from pandas.tslib import iNaT
from six import (
iteritems,
@@ -746,7 +745,7 @@ class BcolzDailyBarReader(DailyBarReader):
return price
class PanelDailyBarReader(DailyBarReader):
class PanelBarReader(DailyBarReader):
"""
Reader for data passed as Panel.
@@ -770,46 +769,54 @@ class PanelDailyBarReader(DailyBarReader):
The first trading day in the dataset.
"""
@preprocess(panel=call(verify_indices_all_unique))
def __init__(self, calendar, panel):
@expect_element(data_frequency={'daily', 'minute'})
def __init__(self, trading_calendar, panel, data_frequency):
panel = panel.copy()
if 'volume' not in panel.minor_axis:
# Fake volume if it does not exist.
panel.loc[:, :, 'volume'] = int(1e9)
self.first_trading_day = panel.major_axis[0]
self._calendar = calendar
self.trading_calendar = trading_calendar
self.first_trading_day = trading_calendar.minute_to_session_label(
panel.major_axis[0]
)
last_trading_day = trading_calendar.minute_to_session_label(
panel.major_axis[-1]
)
self.sessions = trading_calendar.sessions_in_range(
self.first_trading_day,
last_trading_day
)
if data_frequency == 'daily':
self._calendar = self.sessions
elif data_frequency == 'minute':
self._calendar = trading_calendar.minutes_for_sessions_in_range(
self.first_trading_day,
last_trading_day
)
self.panel = panel
@property
def sessions(self):
return self._calendar
sessions = None
@property
def last_available_dt(self):
return self._calendar[-1]
@property
def trading_calendar(self):
return None
trading_calendar = None
def load_raw_arrays(self, columns, start_date, end_date, assets):
columns = list(columns)
def load_raw_arrays(self, columns, start_dt, end_dt, assets):
cal = self._calendar
index = cal[cal.slice_indexer(start_date, end_date)]
shape = (len(index), len(assets))
results = []
for col in columns:
outbuf = zeros(shape=shape)
for i, asset in enumerate(assets):
data = self.panel.loc[asset, start_date:end_date, col]
data = data.reindex_axis(index).values
outbuf[:, i] = data
results.append(outbuf)
return results
return self.panel.loc[
list(assets),
start_dt:end_dt,
list(columns)
].reindex(major_axis=cal[cal.slice_indexer(start_dt, end_dt)]).values.T
def spot_price(self, sid, day, colname):
def spot_price(self, sid, dt, colname):
"""
Parameters
----------
@@ -829,7 +836,9 @@ class PanelDailyBarReader(DailyBarReader):
Returns -1 if the day is within the date range, but the price is
0.
"""
return self.panel.loc[sid, day, colname]
return self.panel.loc[sid, dt, colname]
get_value = spot_price
def get_last_traded_dt(self, sid, dt):
"""
@@ -845,12 +854,9 @@ class PanelDailyBarReader(DailyBarReader):
pd.Timestamp : The last know dt for the asset and dt;
NaT if no trade is found before the given dt.
"""
while dt in self.panel.major_axis:
freq = self.panel.major_axis.freq
if not isnull(self.panel.loc[sid, dt, 'close']):
return dt
dt -= freq
else:
try:
return self.panel.loc[sid, :dt, 'close'].last_valid_index()
except IndexError:
return NaT
+12 -4
View File
@@ -21,6 +21,7 @@ from zipline.finance.trading import TradingEnvironment
from zipline.pipeline.data import USEquityPricing
from zipline.pipeline.loaders import USEquityPricingLoader
from zipline.utils.calendars import get_calendar
from zipline.utils.factory import create_simulation_parameters
import zipline.utils.paths as pth
@@ -150,14 +151,21 @@ def _run(handle_data,
raise ValueError(
"No PipelineLoader registered for column %s." % column
)
else:
env = None
choose_loader = None
perf = TradingAlgorithm(
namespace=namespace,
capital_base=capital_base,
start=start,
end=end,
env=env,
get_pipeline_loader=choose_loader,
sim_params=create_simulation_parameters(
start=start,
end=end,
capital_base=capital_base,
data_frequency=data_frequency,
),
**{
'initialize': initialize,
'handle_data': handle_data,
@@ -314,8 +322,8 @@ def run_algorithm(start,
load_extensions(default_extension, extensions, strict_extensions, environ)
non_none_data = valfilter(bool, {
'data': data,
'bundle': bundle,
'data': data is not None,
'bundle': bundle is not None,
})
if not non_none_data:
# if neither data nor bundle are passed use 'quantopian-quandl'