From 4c5cb867dbd93c45a97aab52ac3a78872c3f03b5 Mon Sep 17 00:00:00 2001
From: Dale Jung <dale@dalejung.com>
Date: Wed, 25 Feb 2015 09:25:35 -0500
Subject: [PATCH] PRF: Sped up the SIDData transforms by using raw values. Also
 fixed a vwap zero division error.

---
 etc/requirements.txt        |  4 ++++
 zipline/protocol.py         | 35 ++++++++++++++++++++++++++---------
 zipline/utils/math_utils.py | 13 ++++++-------
 3 files changed, 36 insertions(+), 16 deletions(-)

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..a41bb459 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):
         """
@@ -360,17 +367,27 @@ class SIDData(object):
         return self._get_bars(days)
 
     def mavg(self, days):
-        return self._get_buffer(self._get_bars(days)).mean()
+        bars = self._get_bars(days)
+        prices = self._get_buffer(bars, raw=True)
+        return nanmean(prices)
 
     def stddev(self, days):
-        return self._get_buffer(self._get_bars(days)).std(ddof=1)
+        bars = self._get_bars(days)
+        prices = self._get_buffer(bars, raw=True)
+        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')
+        prices = self._get_buffer(bars, raw=True)
+        vols = self._get_buffer(bars, field='volume', raw=True)
 
-        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 43e4d380..e1b19bee 100644
--- a/zipline/utils/math_utils.py
+++ b/zipline/utils/math_utils.py
@@ -14,22 +14,21 @@
 # limitations under the License.
 
 import math
-import numpy as np
 
 
 def tolerant_equals(a, b, atol=10e-7, rtol=10e-7):
     return math.fabs(a - b) <= (atol + rtol * math.fabs(b))
 
 
-nanmean = np.nanmean
-nanstd = np.nanstd
-nansum = np.nansum
-
-
 try:
+    # fast versions
     import bottleneck as bn
     nanmean = bn.nanmean
     nanstd = bn.nanstd
     nansum = bn.nansum
 except ImportError:
-    pass
+    # slower numpy
+    import numpy as np
+    nanmean = np.nanmean
+    nanstd = np.nanstd
+    nansum = np.nansum