Merge pull request #43 from quantopian/fawce_sprint6

Fawce sprint6

zipline/finance/movingaverage.py

@@ -0,0 +1,67 @@
+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.average
+        return self.state
+    
+    def create_vwap(self):
+        return DailyVWAP(self.daycount)
+
+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):
+    """
+    Tracks a window of the event history. Use an instance to track the events
+    inside your window to efficiently calculate rolling statistics.
+    """
+    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:]
+
+        

zipline/finance/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/vwap.py

@@ -1,67 +1,61 @@
+import pandas
 from datetime import timedelta
-from itertools import ifilter
 from collections import defaultdict
 
 from zipline.messaging import BaseTransform
+from zipline.finance.movingaverage import EventWindow
 
 class VWAPTransform(BaseTransform):
     
     def init(self, daycount=3):
         self.daycount = daycount
-        self.by_sid = defaultdict(DailyVWAP)
+        self.by_sid = defaultdict(self.create_vwap)
         
     def transform(self, event):
         cur = self.by_sid(event.sid)
         cur.update(event)
         self.state['value'] = cur.vwap
         return self.state
-        
+    
+    def create_vwap(self):
+        return DailyVWAP(self.daycount)
+
 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 = []
+    def __init__(self, daycount):
+        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

zipline/lines.py

@@ -143,6 +143,8 @@ class SimulatedTrading(object):
             sockets[7],
             logging = qutil.LOGGER
         )
+        
+        self.con.cancel_socket = self.allocator.lease(1)[0]
 
         # TODO: Not freeform
         self.con.manage(

zipline/test/test_transforms.py

@@ -0,0 +1,97 @@
+from datetime import timedelta
+from collections import defaultdict
+from unittest2 import TestCase
+
+import zipline.test.factory as factory
+import zipline.util as qutil
+from zipline.finance.vwap import DailyVWAP, VWAPTransform
+from zipline.finance.returns import ReturnsFromPriorClose
+from zipline.finance.movingaverage import MovingAverage
+from zipline.lines import SimulatedTrading
+from zipline.simulator import AddressAllocator, Simulator
+
+
+allocator = AddressAllocator(1000)
+
+class ZiplineWithTransformsTestCase(TestCase):
+    leased_sockets = defaultdict(list)
+    
+    def setUp(self):
+        # skip ahead 100 spots
+        allocator.lease(100)
+        qutil.configure_logging()
+        self.trading_environment = factory.create_trading_environment()
+        self.zipline_test_config = {
+            'allocator':allocator,
+            'sid':133
+        }
+        
+    def test_vwap_tnfm(self):
+        zipline = SimulatedTrading.create_test_zipline(
+            **self.zipline_test_config
+        )
+        
+        vwap = VWAPTransform("vwap_10", daycount=10)
+        zipline.add_transform(vwap)
+        
+        zipline.simulate(blocking=True)
+
+        self.assertTrue(zipline.sim.ready())
+        self.assertFalse(zipline.sim.exception)
+        
+class FinanceTransformsTestCase(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)
+        
+        
+