From 3208bfdbb632d9f6ac5340c4ec8549139c7fa842 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 8 Mar 2016 13:54:46 -0500 Subject: [PATCH] ENH: Add min date to TradingEnvironment. Allow creation of TradingEnvironment to specify a minimum date, so that trading days, market opens, etc. can trimmed to a range more relevant to the backtest. This changes is with an eye towards storing all market minutes in the trading environment, where storing values for much more than the simulation range starts to become more costly. --- tests/test_finance.py | 11 ++++++++++- zipline/finance/trading.py | 6 ++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/test_finance.py b/tests/test_finance.py index 4197de07..ea3a3ce1 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -25,6 +25,7 @@ from unittest import TestCase from datetime import datetime, timedelta import numpy as np +import pandas as pd from nose.tools import timed @@ -546,8 +547,16 @@ class TradingEnvironmentTestCase(TestCase): self.assertTrue(all(friday == minutes[31:421])) self.assertTrue(all(thursday == minutes[421:])) + def test_min_date(self): + min_date = pd.Timestamp('2016-03-04', tz='UTC') + env = TradingEnvironment(min_date=min_date) + + self.assertGreaterEqual(env.first_trading_day, min_date) + self.assertGreaterEqual(env.treasury_curves.index[0], + min_date) + def test_max_date(self): - max_date = datetime(2008, 8, 1, tzinfo=pytz.utc) + max_date = pd.Timestamp('2008-08-01', tz='UTC') env = TradingEnvironment(max_date=max_date) self.assertLessEqual(env.last_trading_day, max_date) diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index d914c611..362d5534 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -68,6 +68,7 @@ class TradingEnvironment(object): load=None, bm_symbol='^GSPC', exchange_tz="US/Eastern", + min_date=None, max_date=None, env_trading_calendar=tradingcalendar, asset_db_path=':memory:' @@ -82,10 +83,7 @@ class TradingEnvironment(object): # `tc_td` is short for "trading calendar trading days" tc_td = env_trading_calendar.trading_days - if max_date: - self.trading_days = tc_td[tc_td <= max_date].copy() - else: - self.trading_days = tc_td.copy() + self.trading_days = tc_td[tc_td.slice_indexer(min_date, max_date)] self.first_trading_day = self.trading_days[0] self.last_trading_day = self.trading_days[-1]