Refactored stdev transform to be applied to multiple fields, not just price

UnitTest for transforms refactored to account for changed stdev transform, introduced AssertAlmostEqual

tests/test_transforms.py

@@ -13,12 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import deque
-
 import pytz
 import numpy as np
 import pandas as pd
 
+from collections import deque
 from datetime import timedelta, datetime
 from unittest import TestCase
 
@@ -273,44 +272,57 @@ class TestFinanceTransforms(TestCase):
         assert tnfm_volumes == expected_volumes
 
     def test_moving_stddev(self):
+
+        stddev = MovingStandardDev(
+            fields=['price', 'volume'],
+            market_aware=False,
+            delta=timedelta(days=3),
+        )
+
         trade_history = factory.create_trade_history(
             133,
             [10.0, 15.0, 13.0, 12.0],
-            [100, 100, 100, 100],
-            timedelta(hours=1),
+            [100, 200, 100, 200],
+            timedelta(days=1),
             self.trading_environment
         )
 
-        stddev = MovingStandardDev(
-            market_aware=False,
-            delta=timedelta(minutes=150),
-        )
         self.source = SpecificEquityTrades(event_list=trade_history)
 
         transformed = list(stddev.transform(self.source))
+        # Output values
+        tnfm_prices = [message.tnfm_value.price for message in transformed]
+        tnfm_volumes = [message.tnfm_value.volume for message in transformed]
 
-        vals = [message.tnfm_value for message in transformed]
-
-        expected = [
+        expected_prices = [
             None,
             np.std([10.0, 15.0], ddof=1),
             np.std([10.0, 15.0, 13.0], ddof=1),
             np.std([15.0, 13.0, 12.0], ddof=1),
         ]
 
-        # np has odd rounding behavior, cf.
-        # http://docs.scipy.org/doc/np/reference/generated/np.std.html
-        for v1, v2 in zip(vals, expected):
+        expected_volumes = [
+            None,
+            np.std([100, 200], ddof=1),
+            np.std([100, 200, 100], ddof=1),
+            np.std([200, 100, 200], ddof=1),
+        ]
 
+        for v1, v2 in zip(tnfm_prices, expected_prices):
             if v1 is None:
                 assert v2 is None
                 continue
-            assert round(v1, 5) == round(v2, 5)
-
+            self.assertAlmostEqual(v1, v2)
 
+        for v1, v2 in zip(tnfm_volumes, expected_volumes):
+            if v1 is None:
+                assert v2 is None
+                continue
+            self.assertAlmostEqual(v1, v2)
 ############################################################
 # Test BatchTransform
 
+
 class TestBatchTransform(TestCase):
     def setUp(self):
         setup_logger(self)

zipline/transforms/stddev.py

@@ -17,19 +17,22 @@ from numbers import Number
 from collections import defaultdict
 from math import sqrt
 
+from zipline import ndict
 from zipline.transforms.utils import EventWindow, TransformMeta
 
 
 class MovingStandardDev(object):
     """
-    Class that maintains a dicitonary from sids to
-    MovingStandardDevWindows.  For each sid, we maintain a the
-    standard deviation of all events falling within the specified
-    window.
+    Class that maintains a dictionary from sids to
+    MovingStandardDevWindows.  For each sid, we maintain standard 
+    deviations over any number of distinct fields. (For example, we can
+    maintain a sid's moving standard deviation of returns as well as its
+    moving standard deviation of prices.
     """
     __metaclass__ = TransformMeta
 
-    def __init__(self, market_aware=True, window_length=None, delta=None):
+    def __init__(self, fields,
+                 market_aware=True, window_length=None, delta=None):
 
         self.market_aware = market_aware
 
@@ -38,6 +41,9 @@ class MovingStandardDev(object):
 
         # Market-aware mode only works with full-day windows.
         if self.market_aware:
+            # Window length must be 1 or greater
+            assert self.window_length >= 1
+
             assert self.window_length and not self.delta,\
                 "Market-aware mode only works with full-day windows."
 
@@ -55,6 +61,7 @@ class MovingStandardDev(object):
         Factory method for self.sid_windows.
         """
         return MovingStandardDevWindow(
+            self.fields,
             self.market_aware,
             self.window_length,
             self.delta
@@ -63,51 +70,86 @@ class MovingStandardDev(object):
     def update(self, event):
         """
         Update the event window for this event's sid.  Return an ndict
-        from tracked fields to moving averages.
+        from tracked fields to moving standard deviations.
         """
         # This will create a new EventWindow if this is the first
         # message for this sid.
         window = self.sid_windows[event.sid]
         window.update(event)
-        return window.get_stddev()
+        return window.get_stddevs()
 
 
 class MovingStandardDevWindow(EventWindow):
     """
-    Iteratively calculates standard deviation for a particular sid
-    over a given time window.  The expected functionality of this
-    class is to be instantiated inside a MovingStandardDev.
+    Iteratively calculates moving standard deviations for a particular sid
+    over a given time window. We can maintain standard deviations for
+    arbitrarily many fields on a single sid. (For example, we might track
+    moving standard deviation of returns as well as its moving standard
+    deviation of prices.) The expected functionality of this class is to be
+    instantiated inside a MovingStandardDev.
     """
 
-    def __init__(self, market_aware, days, delta):
+    def __init__(self, fields, market_aware, window_length, delta):
         # Call the superclass constructor to set up base EventWindow
         # infrastructure.
-        EventWindow.__init__(self, market_aware, days, delta)
+        EventWindow.__init__(self, market_aware, window_length, delta)
 
-        self.sum = 0.0
-        self.sum_sqr = 0.0
+        self.fields = fields
+        self.sum = defaultdict(float)
+        self.sum_sqr = defaultdict(float)
 
     def handle_add(self, event):
-        assert isinstance(event.price, Number)
+        # Sanity check on the event.
+        self.assert_required_fields(event)
 
-        self.sum += event.price
-        self.sum_sqr += event.price ** 2
+        # Increment our running totals with data from the event.
+        for field in self.fields:
+            self.sum[field] += event[field]
+            self.sum_sqr[field] += event[field] ** 2
 
     def handle_remove(self, event):
-        assert isinstance(event.price, Number)
+        # Sanity check on the event.
+        self.assert_required_fields(event)
 
-        self.sum -= event.price
-        self.sum_sqr -= event.price ** 2
+        # Decrement our running totals with data from the event.
+        for field in self.fields:
+            self.sum[field] -= event[field]
+            self.sum_sqr[field] -= event[field] ** 2
 
-    def get_stddev(self):
-        # Sample standard deviation is undefined for a single event or
-        # no events.
-        if len(self) <= 1:
+    def stdev(self, field):
+        """
+        Calculate the standard deviation of our ticks over a single field
+        using a naive algorithm (see http://goo.gl/wPFtf).
+        """
+        # Sanity check.
+        assert field in self.fields
+        # Standard deviation is undefined for no event and 0 for one event
+        if len(self.ticks) <= 1:
             return None
 
+        # Calculate and return the standard deviation.
         else:
-            average = self.sum / len(self)
-            s_squared = (self.sum_sqr - self.sum * average) \
-                / (len(self) - 1)
-            stddev = sqrt(s_squared)
-        return stddev
+            _mean = self.sum[field] / len(self.ticks)
+            _var = (self.sum_sqr[field] -
+                    self.sum[field] * _mean) / (len(self.ticks) - 1)
+            return sqrt(_var)
+
+    def get_stddevs(self):
+        """
+        Return an ndict of all our tracked standard deviations.
+        """
+        out = ndict()
+        for field in self.fields:
+            out[field] = self.stdev(field)
+        return out
+
+    def assert_required_fields(self, event):
+        """
+        We only allow events with all of our tracked fields.
+        """
+        for field in self.fields:
+            assert field in event, \
+                "Event missing [%s] in MovingStandardDevEventWindow" % field
+            assert isinstance(event[field], Number), \
+                "Got %s for %s in MovingStandardDevEventWindow" \
+                % (event[field], field)