mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-07 13:35:43 +08:00
Merge pull request #1313 from nathanwolfe/master
BUG: Add support for Panel data in accordance with documentation
This commit is contained in:
+85
-1
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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'
|
||||
|
||||
Reference in New Issue
Block a user