From 1295f45e13e260f036ab790613347ce8c13486c8 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 13 Aug 2013 11:07:31 -0400 Subject: [PATCH] MAINT: Switch treasury curves from Series to DataFrame. Instead of using a pandas Series of with dictionaries as the values treasury curves, use a DataFrame which more naturally fits the data type of a having a timeseries with mulitple values. Should allow easier slicing/manipulation of the treasury curves, e.g. getting 10 year curves would now be: ``` treasury_curves['10year'] ``` --- zipline/finance/risk/risk.py | 4 ++-- zipline/finance/trading.py | 4 ++-- zipline/utils/factory.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/zipline/finance/risk/risk.py b/zipline/finance/risk/risk.py index c4f22e55..c270fb8e 100644 --- a/zipline/finance/risk/risk.py +++ b/zipline/finance/risk/risk.py @@ -187,7 +187,7 @@ def alpha(algorithm_period_return, treasury_period_return, def get_treasury_rate(treasury_curves, treasury_duration, day): rate = None - curve = treasury_curves[day] + curve = treasury_curves.ix[day] # 1month note data begins in 8/2001, # so we can use 3month instead. idx = TREASURY_DURATIONS.index(treasury_duration) @@ -238,7 +238,7 @@ def choose_treasury(treasury_curves, start_date, end_date): end_day = end_date.replace(hour=0, minute=0, second=0, microsecond=0) search_day = None - if end_day in treasury_curves: + if end_day in treasury_curves.index: rate = get_treasury_rate(treasury_curves, treasury_duration, end_day) diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 264d9a5f..fd492b1b 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -85,9 +85,9 @@ class TradingEnvironment(object): self.benchmark_returns, treasury_curves_map = \ load(self.bm_symbol) - self.treasury_curves = pd.Series(treasury_curves_map) + self.treasury_curves = pd.DataFrame(treasury_curves_map).T if max_date: - self.treasury_curves = self.treasury_curves[:max_date] + self.treasury_curves = self.treasury_curves.ix[:max_date, :] self.full_trading_day = datetime.timedelta(hours=6, minutes=30) self.early_close_trading_day = datetime.timedelta(hours=3, minutes=30) diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index c513bce7..123981b8 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -104,7 +104,7 @@ def create_random_simulation_parameters(): len(treasury_curves) - 1 ) - start_dt = treasury_curves.keys()[random_index] + start_dt = treasury_curves.index[random_index] end_dt = start_dt + timedelta(days=365) now = datetime.utcnow().replace(tzinfo=pytz.utc)