Merge pull request #1313 from nathanwolfe/master

BUG: Add support for Panel data in accordance with documentation

tests/test_algorithm.py

@@ -33,7 +33,10 @@
 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 @@
     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 @@ def analyze(context, results):
                     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()

tests/test_panel_daily_bar_reader.py → tests/test_panel_bar_reader.py

@@ -18,31 +18,29 @@
 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 @@ def init_class_fixtures(cls):
             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 @@ def test_duplicate_values(self):
         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 @@ def test_duplicate_values(self):
     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

zipline/algorithm.py

@@ -37,7 +37,7 @@
 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 @@ def run(self, data=None, overwrite_sim_params=True):
                 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 @@ def run(self, data=None, overwrite_sim_params=True):
                         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

zipline/data/data_portal.py

@@ -156,8 +156,6 @@ def __init__(self,
                 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