ENH: A TradingAlgorithm method called before each trading day

tests/test_tradesimulation.py

@@ -12,12 +12,28 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pandas as pd
 
+from nose_parameterized import parameterized
+from six.moves import range
 from unittest import TestCase
+from zipline import TradingAlgorithm
 from zipline.test_algorithms import NoopAlgorithm
 from zipline.utils import factory
 
 
+class BeforeTradingAlgorithm(TradingAlgorithm):
+    def __init__(self, *args, **kwargs):
+        self.before_trading_at = []
+        super(BeforeTradingAlgorithm, self).__init__(*args, **kwargs)
+
+    def before_trading_start(self):
+        self.before_trading_at.append(self.datetime)
+
+
+FREQUENCIES = {'daily': 0, 'minute': 1}  # daily is less frequent than minute
+
+
 class TestTradeSimulation(TestCase):
 
     def test_minutely_emissions_generate_performance_stats_for_last_day(self):
@@ -27,3 +43,24 @@ class TestTradeSimulation(TestCase):
         algo = NoopAlgorithm(sim_params=params)
         algo.run(source=[])
         self.assertEqual(algo.perf_tracker.day_count, 1.0)
+
+    @parameterized.expand([('%s_%s_%s' % (num_days, freq, emission_rate),
+                            num_days, freq, emission_rate)
+                           for freq in FREQUENCIES
+                           for emission_rate in FREQUENCIES
+                           for num_days in range(1, 4)
+                           if FREQUENCIES[emission_rate] <= FREQUENCIES[freq]])
+    def test_before_trading_start(self, test_name, num_days, freq,
+                                  emission_rate):
+        params = factory.create_simulation_parameters(
+            num_days=num_days, data_frequency=freq,
+            emission_rate=emission_rate)
+
+        algo = BeforeTradingAlgorithm(sim_params=params)
+        algo.run(source=[])
+
+        self.assertEqual(algo.perf_tracker.day_count, num_days)
+        self.assertTrue(params.trading_days.equals(
+            pd.DatetimeIndex(algo.before_trading_at)),
+            "Expected %s but was %s."
+            % (params.trading_days, algo.before_trading_at))

zipline/algorithm.py

@@ -178,25 +178,30 @@ class TradingAlgorithm(object):
         self.algoscript = kwargs.pop('script', None)
 
         self._initialize = None
+        self._before_trading_start = None
         self._analyze = None
 
         if self.algoscript is not None:
             exec_(self.algoscript, self.namespace)
-            self._initialize = self.namespace.get('initialize', None)
+            self._initialize = self.namespace.get('initialize')
             if 'handle_data' not in self.namespace:
                 raise ValueError('You must define a handle_data function.')
             else:
                 self._handle_data = self.namespace['handle_data']
 
+            self._before_trading_start = \
+                self.namespace.get('before_trading_start')
             # Optional analyze function, gets called after run
-            self._analyze = self.namespace.get('analyze', None)
+            self._analyze = self.namespace.get('analyze')
 
-        elif kwargs.get('initialize', False) and kwargs.get('handle_data'):
+        elif kwargs.get('initialize') and kwargs.get('handle_data'):
             if self.algoscript is not None:
                 raise ValueError('You can not set script and \
                 initialize/handle_data.')
             self._initialize = kwargs.pop('initialize')
             self._handle_data = kwargs.pop('handle_data')
+            self._before_trading_start = kwargs.pop('before_trading_start',
+                                                    None)
 
         # If method not defined, NOOP
         if self._initialize is None:
@@ -223,6 +228,12 @@ class TradingAlgorithm(object):
         with ZiplineAPI(self):
             self._initialize(self)
 
+    def before_trading_start(self):
+        if self._before_trading_start is None:
+            return
+
+        self._before_trading_start(self)
+
     def handle_data(self, data):
         if self.history_container:
             self.history_container.update(data, self.datetime)

zipline/gens/tradesimulation.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from logbook import Logger, Processor
+from pandas.tslib import normalize_date
 
 from zipline.finance import trading
 from zipline.protocol import (
@@ -51,10 +52,7 @@ class AlgorithmSimulator(object):
         # Algo Setup
         # ==============
         self.algo = algo
-        self.algo_start = self.sim_params.first_open
-        self.algo_start = self.algo_start.replace(hour=0, minute=0,
-                                                  second=0,
-                                                  microsecond=0)
+        self.algo_start = normalize_date(self.sim_params.first_open)
 
         # ==============
         # Snapshot Setup
@@ -103,10 +101,14 @@ class AlgorithmSimulator(object):
         # inject the current algo
         # snapshot time to any log record generated.
         with self.processor.threadbound():
+            data_frequency = self.sim_params.data_frequency
+
+            self._call_before_trading_start(mkt_open)
+
             for date, snapshot in stream_in:
 
                 self.simulation_dt = date
-                self.algo.on_dt_changed(date)
+                self.on_dt_changed(date)
 
                 # If we're still in the warmup period.  Use the event to
                 # update our universe, but don't yield any perf messages,
@@ -116,8 +118,8 @@ class AlgorithmSimulator(object):
                         if event.type == DATASOURCE_TYPE.SPLIT:
                             self.algo.blotter.process_split(event)
 
-                        if event.type in (DATASOURCE_TYPE.TRADE,
-                                          DATASOURCE_TYPE.CUSTOM):
+                        elif event.type in (DATASOURCE_TYPE.TRADE,
+                                            DATASOURCE_TYPE.CUSTOM):
                             self.update_universe(event)
                         self.algo.perf_tracker.process_event(event)
                 else:
@@ -133,8 +135,8 @@ class AlgorithmSimulator(object):
 
                     # When emitting minutely, we re-iterate the day as a
                     # packet with the entire days performance rolled up.
-                    if self.algo.perf_tracker.emission_rate == 'minute':
-                        if date == mkt_close:
+                    if date == mkt_close:
+                        if self.algo.perf_tracker.emission_rate == 'minute':
                             daily_rollup = self.algo.perf_tracker.to_dict(
                                 emission_type='daily'
                             )
@@ -143,21 +145,37 @@ class AlgorithmSimulator(object):
                             yield daily_rollup
                             tp = self.algo.perf_tracker.todays_performance
                             tp.rollover()
-                            if mkt_close <= self.algo.perf_tracker.last_close:
-                                try:
-                                    mkt_open, mkt_close = \
-                                        trading.environment \
-                                               .next_open_and_close(mkt_close)
 
-                                except trading.NoFurtherDataError:
-                                    # If at the end of backtest history,
-                                    # skip advancing market close.
-                                    pass
+                        if mkt_close <= self.algo.perf_tracker.last_close:
+                            before_last_close = \
+                                mkt_close < self.algo.perf_tracker.last_close
+                            try:
+                                mkt_open, mkt_close = \
+                                    trading.environment \
+                                           .next_open_and_close(mkt_close)
+
+                            except trading.NoFurtherDataError:
+                                # If at the end of backtest history,
+                                # skip advancing market close.
+                                pass
+                            if (self.algo.perf_tracker.emission_rate
+                                    == 'minute'):
                                 self.algo.perf_tracker\
                                          .handle_intraday_market_close(
                                              mkt_open,
                                              mkt_close)
 
+                            if before_last_close:
+                                self._call_before_trading_start(mkt_open)
+
+                    elif data_frequency == 'daily':
+                        next_day = trading.environment.next_trading_day(date)
+
+                        if (next_day is not None
+                                and next_day
+                                < self.algo.perf_tracker.last_close):
+                            self._call_before_trading_start(next_day)
+
                     self.algo.portfolio_needs_update = True
 
             risk_message = self.algo.perf_tracker.handle_simulation_end()
@@ -237,6 +255,16 @@ class AlgorithmSimulator(object):
         self.algo.blotter.new_orders = []
         return orders
 
+    def _call_before_trading_start(self, dt):
+        dt = normalize_date(dt)
+        self.simulation_dt = dt
+        self.on_dt_changed(dt)
+        self.algo.before_trading_start()
+
+    def on_dt_changed(self, dt):
+        if self.algo.datetime != dt:
+            self.algo.on_dt_changed(dt)
+
     def get_message(self, dt):
         """
         Get a perf message for the given datetime.