new-style returns

tests/test_transforms.py

@@ -10,41 +10,12 @@ from zipline.gens.tradegens import SpecificEquityTrades
 from zipline.gens.transform import StatefulTransform
 from zipline.gens.vwap import VWAP
 from zipline.gens.mavg import MovingAverage
-from zipline.finance.returns import ReturnsFromPriorClose
+from zipline.gens.returns import Returns
 from zipline.lines import SimulatedTrading
 from zipline.core.devsimulator import AddressAllocator
 
 allocator = AddressAllocator(1000)
 
-class ZiplineWithTransformsTestCase(TestCase):
-    leased_sockets = defaultdict(list)
-
-    def setUp(self):
-        # skip ahead 100 spots
-        allocator.lease(100)
-        self.trading_environment = factory.create_trading_environment()
-        self.zipline_test_config = {
-            'allocator' : allocator,
-            'sid'       : 133,
-            'devel'     : True
-        }
-        setup_logger(self, '/var/log/qexec/qexec.log')
-
-    def tearDown(self):
-        teardown_logger(self)
-
-    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):
@@ -64,7 +35,12 @@ class FinanceTransformsTestCase(TestCase):
         self.log_handler.pop_application()
 
     def test_vwap(self):
-        vwap = StatefulTransform(VWAP, timedelta(days = 2))
+        
+        vwap = StatefulTransform(
+            VWAP, 
+            market_aware = False,
+            delta = timedelta(days = 2)
+        )
         transformed = list(vwap.transform(self.source))
 
         # Output values
@@ -72,35 +48,32 @@ class FinanceTransformsTestCase(TestCase):
         # "Hand calculated" values.
         expected = [(10.0 * 100) / 100.0,
                     ((10.0 * 100) + (10.0 * 100)) / (200.0),
-                    ((10.0 * 100) + (10.0 * 100) + (11.0 * 100)) / (300.0),
-                    # First event should get droppped here.
-                    ((10.0 * 100) + (11.0 * 100) + (11.0 * 300)) / (500.0)] 
+                    # We should drop the first event here.
+                    ((10.0 * 100) + (11.0 * 100)) / (200.0),
+                    # We should drop the second event here.
+                    ((11.0 * 100) + (11.0 * 300)) / (400.0)] 
 
         # Output should match the expected.
         assert tnfm_vals == expected
         
-
     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()
+        returns = StatefulTransform(
+            Returns
         for trade in trade_history:
             returns.update(trade)
-
-
         self.assertEqual(returns.returns, .1)
 
-
     def test_moving_average(self):
-        
+
         mavg = StatefulTransform(
             MovingAverage,          
+            market_aware = False,
             fields = ['price', 'volume'],
             delta = timedelta(days = 2), 
         ) 
@@ -109,17 +82,21 @@ class FinanceTransformsTestCase(TestCase):
         # Output values.
         tnfm_prices = [message.tnfm_value.price for message in transformed]
         tnfm_volumes = [message.tnfm_value.volume for message in transformed]
+
         # "Hand-calculated" values
         expected_prices = [((10.0) / 1.0),
                            ((10.0 + 10.0) / 2.0),
-                           ((10.0 + 10.0 + 11.0) / 3.0),
                            # First event should get dropped here.
-                           ((10.0 + 11.0 + 11.0) / 3.0)]
+                           ((10.0 + 11.0) / 2.0),
+                           # Second event should get dropped here.
+                           ((11.0 + 11.0) / 2.0)]
+
         expected_volumes = [((100.0) / 1.0),
                            ((100.0 + 100.0) / 2.0),
-                           ((100.0 + 100.0 + 100.0) / 3.0),
-                           # First event should get dropped here.
-                           ((100.0 + 100.0 + 300.0) / 3.0)]
+                            # First event should get dropped here.
+                           ((100.0 + 100.0) / 2.0),
+                           # Second event should get dropped here.
+                           ((100.0 + 300.0) / 2.0)]
         
         assert tnfm_prices == expected_prices
         assert tnfm_volumes == expected_volumes

zipline/finance/returns.py

@@ -1,47 +0,0 @@
-from collections import defaultdict
-from zipline.transforms.base import BaseTransform
-
-class Returns(object):
-    """
-    Class that maintains a dictionary from sids to the event
-    representing the most recent closing price.
-    """
-    def __init__(self, days == 1):
-        self.days = days
-        self.mapping = defaultdict(self._create)
-        
-    def update(self, event):
-        """
-        Update and return the calculated returns for this event's sid.
-        """
-        sid_returns = self.mapping[event.sid].update(event)
-        return sid_returns
-
-    def _create(self):
-        return ReturnsFromPriorClose(days)
-
-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):
-        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/gens/returns.py

@@ -0,0 +1,77 @@
+from collections import defaultdict
+from zipline.transforms.base import BaseTransform
+from zipline.utils.tradingcalendar import market_closes
+
+class Returns(object):
+    """
+    Class that maintains a dictionary from sids to the sid's
+    closing price N trading days ago.
+    """
+    def __init__(self, days):
+        self.days = days
+        self.mapping = defaultdict(self._create)
+        
+    def update(self, event):
+        """
+        Update and return the calculated returns for this event's sid.
+        """
+        assert event.has_key('dt')
+        assert event.has_key('price')
+
+        tracker = self.mapping[event.sid]
+        tracker.update(event)
+        
+        return tracker.get_returns()
+
+    def _create(self):
+        return ReturnsFromPriorClose(days)
+
+class ReturnsFromPriorClose(object):
+    """
+    Records the last N closing events for a given security as well as the
+    last event for the security.  When we get an event for a new day, we 
+    treat the last event seen  as the close for the previous day.
+    """
+
+    def __init__(self, days):
+        self.closes = deque()
+        self.last_event = None
+        self.returns = None
+        self.days = days
+
+    def get_returns(self):
+        return self.returns
+
+    def update(self, event):
+        
+        if self.last_event:
+
+            # Day has changed since the last event we saw.  Treat
+            # the last event as the closing price for its day and 
+            # clear out the oldest close if it has expired.
+            if self.last_event.dt.date() != event.dt.date():
+
+                self.closes.append(self.last_event)
+
+                # We keep an event for the end of each trading day, so
+                # if the number of stored events is greater than the
+                # number of days we want to track, the oldest close
+                # is expired and should be discarded.
+                if len(self.closes) > self.days:
+                    # Pop the oldest event.
+                    self.closes.popleft()
+
+        # We only generate a return value once we've seen enough days
+        # to give a sensible value.  Would be nice if we could query
+        # db for closes prior to our initial event, but that would
+        # require giving this transform database creds, which we want
+        # to avoid.
+
+        if len(self.closes) == self.days:
+            change = event.price - self.closes[0].price
+            self.returns = change / self.last_close.price
+
+
+        # the current event is now the last_event
+        self.last_event = event
+        

zipline/gens/transform.py

@@ -26,6 +26,7 @@ class Passthrough(object):
     def update(self, event):
         pass
 
+# Deprecated
 def functional_transform(stream_in, func, *args, **kwargs):
     """
     Generic transform generator that takes each message from an in-stream
@@ -213,7 +214,6 @@ class EventWindow:
         #                              oldest               newest
         #                                |                    |
         #                                V                    V
-        import nose.tools; nose.tools.set_trace()
         while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt):
             
             # popleft removes and returns the oldest tick in self.ticks
@@ -229,7 +229,6 @@ class EventWindow:
     def out_of_delta(self, oldest, newest):
         return (newest - oldest) >= self.delta
 
-
     # All event windows expect to receive events with datetime fields
     # that arrive in sorted order.
     def assert_well_formed(self, event):

zipline/utils/tradingcalendar.py

@@ -346,11 +346,13 @@ opens = rrule.rruleset(cache=True)
 opens.rrule(market_opens_with_holidays)
 for holiday_rule in holiday_opens:
     opens.exrule(holiday_rule)
-open_count = opens.count()
 
 closes = rrule.rruleset(cache=True)
 closes.rrule(market_closes_with_holidays)
 for holiday_rule in holiday_closes:
     closes.exrule(holiday_rule)
+
+# This runs the calendar to load all data into a cache.
+open_count = opens.count()
 close_count = closes.count()