diff --git a/etc/requirements.txt b/etc/requirements.txt index fa21ab3d..0d8dcfec 100644 --- a/etc/requirements.txt +++ b/etc/requirements.txt @@ -24,3 +24,7 @@ Cython==0.20.1 # faster OrderedDict cyordereddict==0.2.2 + +# faster array ops. Note once numpy gets to 1.9.1 we +# we can bump to 1.0.0 +bottleneck==0.8.0 diff --git a/zipline/protocol.py b/zipline/protocol.py index e364fa01..0fd5c43e 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -15,8 +15,10 @@ from six import iteritems, iterkeys import pandas as pd +import numpy as np from . utils.protocol_utils import Enum +from . utils.math_utils import nanstd, nanmean, nansum from zipline.finance.trading import with_environment from zipline.utils.algo_instance import get_algo_instance @@ -270,7 +272,7 @@ class SIDData(object): def __repr__(self): return "SIDData({0})".format(self.__dict__) - def _get_buffer(self, bars, field='price'): + def _get_buffer(self, bars, field='price', raw=False): """ Gets the result of history for the given number of bars and field. @@ -287,7 +289,7 @@ class SIDData(object): cls._history_cache = {} if field not in self._history_cache \ - or bars > len(cls._history_cache[field].index): + or bars > len(cls._history_cache[field][0].index): # If we have never cached this field OR the amount of bars that we # need for this field is greater than the amount we have cached, # then we need to get more history. @@ -297,12 +299,17 @@ class SIDData(object): # Assert that the column holds ints, not security objects. if not isinstance(self._sid, str): hst.columns = hst.columns.astype(int) - self._history_cache[field] = hst + self._history_cache[field] = (hst, hst.values, hst.columns) # Slice of only the bars needed. This is because we strore the LARGEST # amount of history for the field, and we might request less than the # largest from the cache. - return cls._history_cache[field][self._sid][-bars:] + buffer_, values, columns = cls._history_cache[field] + if raw: + sid_index = columns.get_loc(self._sid) + return values[-bars:, sid_index] + else: + return buffer_[self._sid][-bars:] def _get_bars(self, days): """ @@ -314,6 +321,14 @@ class SIDData(object): point to a new function object. """ + def daily_get_max_bars(days): + return days + + def minute_get_max_bars(days): + # max number of minute. regardless of current days or short + # sessions + return days * 390 + def daily_get_bars(days): return days @@ -351,26 +366,41 @@ class SIDData(object): self._freqstr = '1d' # update this method to point to the daily variant. self._get_bars = daily_get_bars + self._get_max_bars = daily_get_max_bars else: self._freqstr = '1m' # update this method to point to the minute variant. self._get_bars = minute_get_bars + self._get_max_bars = minute_get_max_bars # Not actually recursive because we have already cached the new method. return self._get_bars(days) def mavg(self, days): - return self._get_buffer(self._get_bars(days)).mean() + bars = self._get_bars(days) + max_bars = self._get_max_bars(days) + prices = self._get_buffer(max_bars, raw=True)[-bars:] + return nanmean(prices) def stddev(self, days): - return self._get_buffer(self._get_bars(days)).std(ddof=1) + bars = self._get_bars(days) + max_bars = self._get_max_bars(days) + prices = self._get_buffer(max_bars, raw=True)[-bars:] + return nanstd(prices, ddof=1) def vwap(self, days): bars = self._get_bars(days) - prices = self._get_buffer(bars) - vols = self._get_buffer(bars, field='volume') + max_bars = self._get_max_bars(days) + prices = self._get_buffer(max_bars, raw=True)[-bars:] + vols = self._get_buffer(max_bars, field='volume', raw=True)[-bars:] - return (prices * vols).sum() / vols.sum() + vol_sum = nansum(vols) + try: + ret = nansum(prices * vols) / vol_sum + except ZeroDivisionError: + ret = np.nan + + return ret def returns(self): algo = get_algo_instance() diff --git a/zipline/utils/math_utils.py b/zipline/utils/math_utils.py index 647ac0b6..e1b19bee 100644 --- a/zipline/utils/math_utils.py +++ b/zipline/utils/math_utils.py @@ -18,3 +18,17 @@ import math def tolerant_equals(a, b, atol=10e-7, rtol=10e-7): return math.fabs(a - b) <= (atol + rtol * math.fabs(b)) + + +try: + # fast versions + import bottleneck as bn + nanmean = bn.nanmean + nanstd = bn.nanstd + nansum = bn.nansum +except ImportError: + # slower numpy + import numpy as np + nanmean = np.nanmean + nanstd = np.nanstd + nansum = np.nansum