refactored transforms to be a package. added vwap, moving average, returns. basic tests for calculations, need to add a test to run the calculations in a zipline.

zipline/finance/transforms/moving_average.py

@@ -0,0 +1,92 @@
+import pandas
+from datetime import timedelta
+from collections import defaultdict
+
+from zipline.messaging import BaseTransform
+
+class MovingAverageTransform(BaseTransform):
+    
+    def init(self, daycount=3):
+        self.daycount = daycount
+        self.by_sid = defaultdict(MovingAverage)
+        
+    def transform(self, event):
+        cur = self.by_sid(event.sid)
+        cur.update(event)
+        self.state['value'] = cur.vwap
+        return self.state
+
+class MovingAverage(object):
+    
+    def __init__(self, daycount):
+        self.window = EventWindow(daycount)
+        self.total = 0.0
+        self.average = 0.0
+        
+    def update(self, event):
+        self.window.update(event)
+        
+        self.total += event.price
+        
+        for dropped in self.window.dropped_ticks:
+            self.total -= dropped.price
+            
+        if len(self.window.ticks) > 0:
+            self.average = self.total / len(self.window.ticks)
+        else:
+            self.average = 0.0
+
+class EventWindow(object):
+    def __init__(self, daycount):
+        self.ticks = []
+        self.dropped_ticks = []
+        self.delta = timedelta(days=daycount)
+    
+    def update(self, event):
+        # add new event
+        self.ticks.append(event) 
+        # determine which events are expired
+        last_date = event['dt']
+        first_date = last_date - self.delta
+        
+        self.dropped_ticks = []
+        for tick in self.ticks:
+            if tick['dt'] <= first_date:
+                self.dropped_ticks.append(tick)
+
+        # remove the expired events
+        slice_index = len(self.dropped_ticks)      
+        self.ticks = self.ticks[slice_index:]
+        
+# ------------------------------
+# Experimental
+# ------------------------------
+      
+class EventHistory(object):
+    
+    def __init__(self, daycount):
+        self.ticks = []
+        self.dropped_ticks = []
+        self.frame = pandas.DataFrame()
+        self.delta = timedelta(days=daycount)
+        
+    def update(self, event):
+        self.ticks.append(event.__dict__)
+        self.last_date = event['dt']
+        self.first_date = self.last_date - self.delta
+        
+        # determine which events are expired
+        self.dropped_ticks = []
+        for tick in self.ticks:
+            if tick['dt'] < self.first_date:
+                self.dropped_ticks.append(tick)
+              
+        # remove the expired events
+        slice_index = len(self.dropped_ticks)      
+        self.ticks = self.ticks[slice_index:]
+        self.frame = pandas.DataFrame(
+            self.ticks
+        )
+        self.frame.index = self.frame['dt']
+
+        

zipline/finance/transforms/returns.py

@@ -0,0 +1,44 @@
+import pandas
+from datetime import timedelta
+from collections import defaultdict
+
+from zipline.messaging import BaseTransform
+
+class WindowTransform(BaseTransform):
+    
+    def init(self, daycount=3):
+        self.daycount = daycount
+        self.by_sid = defaultdict(DailyReturns)
+        
+    def transform(self, event):
+        cur = self.by_sid(event.sid)
+        cur.update(event)
+        self.state['value'] = cur.vwap
+        return self.state
+
+class ReturnsFromPriorClose(object):
+    """
+    Calculates a security's returns since the previous close, using the
+    current price.
+    """
+    
+    def __init__(self):
+        self.last_close = None
+        self.last_event = None
+        self.returns = 0.0
+        
+    def update(self, event):
+        next_close = None
+        if self.last_close:
+            change = event.price - self.last_close.price
+            self.returns = change / self.last_close.price
+            
+        if self.last_event:
+            if self.last_event.dt.day != event.dt.day:
+                # the current event is from the day after
+                # the last event. Therefore the last event was
+                # the last close
+                self.last_close = self.last_event
+        
+        # the current event is now the last_event
+        self.last_event = event

zipline/finance/transforms/vwap.py

@@ -1,8 +1,9 @@
+import pandas
 from datetime import timedelta
-from itertools import ifilter
 from collections import defaultdict
 
 from zipline.messaging import BaseTransform
+from zipline.finance.transforms.moving_average import EventWindow, EventHistory
 
 class VWAPTransform(BaseTransform):
     
@@ -15,53 +16,72 @@ class VWAPTransform(BaseTransform):
         cur.update(event)
         self.state['value'] = cur.vwap
         return self.state
-        
+
 class DailyVWAP:
     """A class that tracks the volume weighted average price
        based on tick updates."""
     def __init__(self, daycount=3):
-        self.ticks = []
-        self.dropped_ticks = []
+        self.window = EventWindow(daycount)
         self.flux = 0.0
         self.volume = 0
-        self.lastTick = None
         self.vwap = 0.0
         self.delta = timedelta(days=daycount)
-    
+
     def update(self, event):
-        
-        self.ticks.append(event)
+
+        # update the event window
+        self.window.update(event)
+
+        # add the current event's flux and volume to the tracker
         flux, volume = self.calculate_flux([event])
         self.flux += flux
         self.volume += volume
-        
-        self.last_date = event['dt']
-        self.first_date = self.last_date - self.delta
-        #use a list comprehension to filter the ticks to those within 
-        #desired day range. The dt properties are full datetime objects
-        #and provide overloads for arithmetic operations. 
-        self.dropped_ticks = []
-        for tick in self.ticks:
-            if tick['dt'] < self.first_date:
-                self.dropped_ticks.append(tick)
-              
-        slice_index = len(self.dropped_ticks)      
-        self.ticks = self.ticks[slice_index:]
 
-        dropped_flux, dropped_volume = self.calculate_flux(self.dropped_ticks)
-        
+        # subract the expired events flux and volume from the tracker
+        dropped = self.window.dropped_ticks
+        dropped_flux, dropped_volume = self.calculate_flux(dropped)
+
         self.flux -= dropped_flux
         self.volume -= dropped_volume
-                                      
+
         if(self.volume != 0):
             self.vwap = self.flux / self.volume
         else:
             self.vwap = None
-            
+
     def calculate_flux(self, ticks):
         flux = 0.0
         volume = 0
         for tick in ticks:
             flux += tick['volume'] * tick['price']
             volume += tick['volume']
-        return flux, volume
+        return flux, volume
+
+
+# ------------------------------
+# Experimental
+# ------------------------------
+        
+class DailyVWAP_df(object):
+    
+    def __init__(self, daycount=3):
+        self.history = EventHistory(daycount)
+        self.vwap = None
+        
+    def update(self, event):
+        self.history.update(event)
+        frame = self.history.frame
+        
+        window = len(frame)
+        value = pandas.rolling_sum(
+            frame['price'] * frame['volume'],
+            window
+        )
+        volume = pandas.rolling_sum(
+            frame['volume'],
+            window
+        )
+        
+        vwap = value / volume
+        self.vwap = vwap[-1]
+        

zipline/test/test_transforms.py

@@ -0,0 +1,61 @@
+from datetime import timedelta
+
+from unittest2 import TestCase
+import zipline.test.factory as factory
+from zipline.finance.transforms.vwap import DailyVWAP, DailyVWAP_df
+from zipline.finance.transforms.returns import ReturnsFromPriorClose
+from zipline.finance.transforms.moving_average import MovingAverage
+
+
+class FinanceTestCase(TestCase):
+    def setUp(self):
+        self.trading_environment = factory.create_trading_environment()
+    def test_vwap(self):
+        
+        trade_history = factory.create_trade_history(
+            133, 
+            [10.0, 10.0, 10.0, 11.0], 
+            [100, 100, 100, 300], 
+            timedelta(days=1), 
+            self.trading_environment
+        )
+        
+        vwap = DailyVWAP(daycount=2)
+        for trade in trade_history:
+            vwap.update(trade)
+        
+        self.assertEqual(vwap.vwap, 10.75)
+        
+    
+    def test_returns(self):
+        trade_history = factory.create_trade_history(
+            133, 
+            [10.0, 10.0, 10.0, 11.0], 
+            [100, 100, 100, 300], 
+            timedelta(days=1), 
+            self.trading_environment
+        )
+        
+        returns = ReturnsFromPriorClose()
+        for trade in trade_history:
+            returns.update(trade)
+            
+        
+        self.assertEqual(returns.returns, .1)
+        
+    
+    def test_moving_average(self):
+        trade_history = factory.create_trade_history(
+            133, 
+            [10.0, 10.0, 10.0, 11.0], 
+            [100, 100, 100, 300], 
+            timedelta(days=1), 
+            self.trading_environment
+        )
+        
+        ma = MovingAverage(daycount=2)
+        for trade in trade_history:
+            ma.update(trade)
+            
+        
+        self.assertEqual(ma.average, 10.5)