Handle missing historical data more elegantly

Updated the search for treasury data when there is none for the
test end date.
It could be that the end date is not a trading day, or we could
just be missing treasury data. In either case, we try to recover
more gracefully now, by searching as far as possible and maybe
logging a warning.

Similarly, if there is no benchmark data for the test end date,
look for the next trading day. If we really have no data,
blow up with our own explicit exception, instead of overflowing
in our search for dates in the future.

zipline/finance/risk.py

@@ -55,6 +55,8 @@ Risk Report
 import logbook
 import datetime
 import math
+import bisect
+from operator import itemgetter
 import numpy as np
 import numpy.linalg as la
 from zipline.utils.date_utils import epoch_now
@@ -321,32 +323,67 @@ class RiskMetricsBase(object):
         else:
             self.treasury_duration = '30year'
 
-        one_day = datetime.timedelta(days=1)
+        # in case end date is not a trading day, search for the next or
+        # previous market day for an interest rate.  choose next in a tie.
+        search_day = None
 
-        curve = None
-        # in case end date is not a trading day, search for the next market
-        # day for an interest rate
-        for i in xrange(7):
-            if (self.end_date + i * one_day) in self.treasury_curves:
-                curve = self.treasury_curves[self.end_date + i * one_day]
-                self.treasury_curve = curve
-                rate = self.treasury_curve[self.treasury_duration]
-                # 1month note data begins in 8/2001,
-                # so we can use 3month instead.
-                if rate is None and self.treasury_duration == '1month':
-                    rate = self.treasury_curve['3month']
+        if self.end_date in self.treasury_curves:
+            search_day = self.end_date
+        else:
+            search_days = self.treasury_curves.keys()
+            next_day = prev_day = None
 
-                if rate is not None:
-                    return rate * (td.days + 1) / 365
+            # Find leftmost item greater than or equal to end_date
+            i = bisect.bisect_left(search_days, self.end_date)
+            if i != len(search_days):
+                next_day = search_days[i]
+            if i:
+                prev_day = search_days[i - 1]
 
-        message = "no rate for end date = {dt} and term = {term}. Check \
-        that date doesn't exceed treasury history range."
+            search_dist = None
+            if next_day and prev_day:
+                search_day, search_dist = \
+                    min(((dt, self.search_day_distance(dt))
+                        for dt in (next_day, prev_day)), key=itemgetter(1))
+            else:
+                search_day = next_day or prev_day
+
+            if search_day:
+                search_dist = search_dist or \
+                    self.search_day_distance(search_day)
+                if search_dist is None or search_dist > 1:
+                    message = "No rate within 1 trading day of end date = \
+{dt} and term = {term}. Check that date doesn't exceed treasury history range."
+                    message = message.format(dt=self.end_date,
+                                             term=self.treasury_duration)
+                    log.warn(message)
+
+        if search_day:
+            curve = self.treasury_curves[search_day]
+            self.treasury_curve = curve
+            rate = self.treasury_curve[self.treasury_duration]
+            # 1month note data begins in 8/2001,
+            # so we can use 3month instead.
+            if rate is None and self.treasury_duration == '1month':
+                rate = self.treasury_curve['3month']
+
+            if rate is not None:
+                return rate * (td.days + 1) / 365
+
+        message = "No rate for end date = {dt} and term = {term}. Check \
+that date doesn't exceed treasury history range."
         message = message.format(
             dt=self.end_date,
             term=self.treasury_duration
         )
         raise Exception(message)
 
+    def search_day_distance(self, dt):
+        tdd = self.trading_environment.trading_day_distance(self.end_date, dt)
+        if tdd is None:
+            return None
+        return tdd if tdd >= 0 else -1 * tdd + .5  # prev is 'farther'
+
 
 class RiskMetricsIterative(RiskMetricsBase):
     """Iterative version of RiskMetrics.
@@ -394,10 +431,16 @@ class RiskMetricsIterative(RiskMetricsBase):
             self.trading_days += 1
             self.update_compounded_log_returns()
 
-        self.end_date += datetime.timedelta(hours=24)
+        next_trading_day = \
+            self.trading_environment.next_trading_day(self.end_date)
 
-        while not self.trading_environment.is_trading_day(self.end_date):
-            self.end_date += datetime.timedelta(hours=24)
+        if next_trading_day:
+            self.end_date = next_trading_day
+        else:
+            message = "No trading data on or after {dt}. Check \
+that date doesn't exceed benchmark history range."
+            message = message.format(dt=self.end_date)
+            raise Exception(message)
 
         self.end_date = self.end_date.replace(hour=0, minute=0, second=0)
 
@@ -408,7 +451,7 @@ class RiskMetricsIterative(RiskMetricsBase):
 
         if(len(self.benchmark_returns) != len(self.algorithm_returns)):
             message = "Mismatch between benchmark_returns ({bm_count}) and \
-            algorithm_returns ({algo_count}) in range {start} : {end}"
+algorithm_returns ({algo_count}) in range {start} : {end}"
             message = message.format(
                 bm_count=len(self.benchmark_returns),
                 algo_count=len(self.algorithm_returns),

zipline/finance/trading.py

@@ -17,7 +17,8 @@ import pytz
 import logbook
 import datetime
 
-from collections import defaultdict
+from collections import defaultdict, OrderedDict
+import bisect
 
 import zipline.protocol as zp
 from zipline.finance.slippage import (
@@ -67,7 +68,7 @@ class TradingEnvironment(object):
     ):
 
         self.trading_days = []
-        self.trading_day_map = {}
+        self.trading_day_map = OrderedDict()
         self.treasury_curves = treasury_curves
         self.benchmark_returns = benchmark_returns
         self.period_start = period_start
@@ -190,6 +191,33 @@ class TradingEnvironment(object):
         dt = self.normalize_date(test_date)
         return (dt in self.trading_day_map)
 
+    def next_trading_day(self, test_date):
+        dt = self.normalize_date(test_date)
+        last_dt = next(reversed(self.trading_day_map))
+        delta = datetime.timedelta(days=1)
+
+        while dt <= last_dt:
+            dt += delta
+            if dt in self.trading_day_map:
+                return dt
+
+        return None
+
+    def trading_day_distance(self, first_date, second_date):
+        first_date = self.normalize_date(first_date)
+        second_date = self.normalize_date(second_date)
+
+        trading_days = self.trading_day_map.keys()
+        # Find leftmost item greater than or equal to day
+        i = bisect.bisect_left(trading_days, first_date)
+        if i == len(trading_days):  # nothing found
+            return None
+        j = bisect.bisect_left(trading_days, second_date)
+        if j == len(trading_days):
+            return None
+
+        return j - i
+
     def get_benchmark_daily_return(self, test_date):
         date = self.normalize_date(test_date)
         if date in self.trading_day_map:

zipline/utils/factory.py

@@ -89,6 +89,10 @@ Fetching data from data.treasury.gov
 
     fp_tr.close()
 
+    tr_curves = OrderedDict(sorted(
+                            ((dt, c) for dt, c in tr_curves.iteritems()),
+                            key=lambda t: t[0]))
+
     return bm_returns, tr_curves