From 3784ed4ba9bf635f29456a58eb70fd4d75d53219 Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Fri, 22 Aug 2014 09:30:57 -0400 Subject: [PATCH] ENH: A TradingAlgorithm method called before each trading day --- tests/test_tradesimulation.py | 37 +++++++++++++++++++ zipline/algorithm.py | 17 +++++++-- zipline/gens/tradesimulation.py | 64 +++++++++++++++++++++++---------- 3 files changed, 97 insertions(+), 21 deletions(-) diff --git a/tests/test_tradesimulation.py b/tests/test_tradesimulation.py index 868d2ae0..f6490cb9 100644 --- a/tests/test_tradesimulation.py +++ b/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)) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 1ccfe46a..244d7552 100644 --- a/zipline/algorithm.py +++ b/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) diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index c465d5b8..9f378869 100644 --- a/zipline/gens/tradesimulation.py +++ b/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.