diff --git a/tests/test_transforms.py b/tests/test_transforms.py index b9420633..2a13e4d7 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -100,9 +100,9 @@ class FinanceTransformsTestCase(TestCase): def test_moving_average(self): mavg = StatefulTransform( - MovingAverage, - timedelta(days = 2), - ['price', 'volume'] + MovingAverage, + fields = ['price', 'volume'], + delta = timedelta(days = 2), ) transformed = list(mavg.transform(self.source)) diff --git a/zipline/gens/mavg.py b/zipline/gens/mavg.py index 1978f92e..21aa0bd0 100644 --- a/zipline/gens/mavg.py +++ b/zipline/gens/mavg.py @@ -13,17 +13,39 @@ class MovingAverage(object): maintain a sid's average volume as well as its average price.) """ - def __init__(self, delta, fields): - self.delta = delta + def __init__(self, fields, market_aware, days = None, delta = None): + self.fields = fields + self.market_aware = market_aware + + self.delta = delta + self.days = days + + # Market-aware mode only works with full-day windows. + if self.market_aware: + assert self.days and not self.delta,\ + "Market-aware mode only works with full-day windows." + + # Non-market-aware mode requires a timedelta. + else: + assert self.delta and not self.days, \ + "Non-market-aware mode requires a timedelta." + # No way to pass arguments to the defaultdict factory, so we # need to define a method to generate the correct EventWindows. self.sid_windows = defaultdict(self.create_window) def create_window(self): - """Factory method for self.sid_windows.""" - return MovingAverageEventWindow(self.delta, self.fields) - + """ + Factory method for self.sid_windows. + """ + return MovingAverageEventWindow( + self.fields, + self.market_aware, + self.days, + self.delta + ) + def update(self, event): """ Update the event window for this event's sid. Return an ndict @@ -45,11 +67,11 @@ class MovingAverageEventWindow(EventWindow): MovingAverage transform. """ - def __init__(self, delta, fields): + def __init__(self, fields, market_aware, days, delta): # Call the superclass constructor to set up base EventWindow # infrastructure. - EventWindow.__init__(self, delta) + EventWindow.__init__(self, market_aware, days, delta) # We maintain a dictionary of totals for each of our tracked # fields. diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 972f44e8..522d80da 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -74,7 +74,7 @@ class TradeSimulationClient(object): # Pipe the events with transactions to perf. This will remove # the txn field added by TransactionSimulator and replace it # with a portfolio object to be passed to the user's - # algorithm. Also adds a PERF_MESSAGE field which is usually + # algorithm. Also adds a perf_message field which is usually # none, but contains an update message once per day. perf_tracker = StatefulTransform( PerformanceTracker, diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 202acc0f..09f074ae 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -2,9 +2,10 @@ Generator versions of transforms. """ import types +import pytz from copy import deepcopy -from datetime import datetime +from datetime import datetime, timedelta from collections import deque, defaultdict from numbers import Number from abc import ABCMeta, abstractmethod @@ -149,6 +150,9 @@ class EventWindow: window. Calls self.handle_remove(event) for each event removed from the window. Subclass these methods along with init(*args, **kwargs) to calculate metrics over the window. + + The market_aware flag is used to toggle whether the eventwindow + calculates See zipline/gens/mavg.py and zipline/gens/vwap.py for example implementations of moving average and volume-weighted average @@ -157,10 +161,31 @@ class EventWindow: # Mark this as an abstract base class. __metaclass__ = ABCMeta - def __init__(self, delta): + def __init__(self, market_aware, days = None, delta = None): + + self.market_aware = market_aware + self.days = days + self.delta = delta + self.ticks = deque() - self.delta = delta + + # Market-aware mode only works with full-day windows. + if self.market_aware: + assert self.days and not self.delta,\ + "Market-aware mode only works with full-day windows." + + # Non-market-aware mode requires a timedelta. + else: + assert self.delta and not self.days, \ + "Non-market-aware mode requires a timedelta." + # Set the behavior for dropping events from the back of the + # event window. + if self.market_aware: + self.drop_condition = self.out_of_market_window + else: + self.drop_condition = self.out_of_timedelta + @abstractmethod def handle_add(self, event): raise NotImplementedError() @@ -174,22 +199,36 @@ class EventWindow: def update(self, event): self.assert_well_formed(event) + # Add new event and increment totals. self.ticks.append(event) + + # Subclasses should override handle_add to define behavior for + # adding new ticks. self.handle_add(event) - # Clear out expired event. + # Clear out any expired events. drop_condition changes depending + # on whether or not we are running in market_aware mode. # - # newest oldest - # | | - # V V - while (self.ticks[-1].dt - self.ticks[0].dt) > self.delta: + # oldest newest + # | | + # V V + while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt): + # popleft removes and returns the oldest tick in self.ticks popped = self.ticks.popleft() + # Subclasses should override handle_remove to define # behavior for removing ticks. self.handle_remove(popped) + def out_of_market_window(self, oldest, newest): + return trading_days_between(oldest, newest) >= self.days + + 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): @@ -200,3 +239,6 @@ class EventWindow: # Something is wrong if new event is older than previous. assert event.dt >= self.ticks[-1].dt, \ "Events arrived out of order in EventWindow: %s -> %s" % (event, self.ticks[0]) + + + diff --git a/zipline/gens/vwap.py b/zipline/gens/vwap.py index 029284c1..5a0947d8 100644 --- a/zipline/gens/vwap.py +++ b/zipline/gens/vwap.py @@ -9,16 +9,33 @@ class VWAP(object): """ Class that maintains a dictionary from sids to VWAPEventWindows. """ - def __init__(self, delta): + def __init__(self, market_aware, delta=None, days=None): + + self.market_aware = market_aware self.delta = delta + self.days = days + + # Market-aware mode only works with full-day windows. + if self.market_aware: + assert self.days and not self.delta,\ + "Market-aware mode only works with full-day windows." + + # Non-market-aware mode requires a timedelta. + else: + assert self.delta and not self.days, \ + "Non-market-aware mode requires a timedelta." # No way to pass arguments to the defaultdict factory, so we # need to define a method to generate the correct EventWindows. self.sid_windows = defaultdict(self.create_window) - + def create_window(self): """Factory method for self.sid_windows.""" - return VWAPEventWindow(self.delta) + return VWAPEventWindow( + self.market_aware, + days = self.days, + delta = self.delta + ) def update(self, event): """ @@ -31,14 +48,13 @@ class VWAP(object): window.update(event) return window.get_vwap() - class VWAPEventWindow(EventWindow): """ Iteratively maintains a vwap for a single sid over a given timedelta. """ - def __init__(self, delta): - EventWindow.__init__(self, delta) + def __init__(self, market_aware, days=None, delta=None): + EventWindow.__init__(self, market_aware, days, delta) self.flux = 0.0 self.totalvolume = 0.0 diff --git a/zipline/utils/tradingcalendar.py b/zipline/utils/tradingcalendar.py index c7aa9152..7742b714 100644 --- a/zipline/utils/tradingcalendar.py +++ b/zipline/utils/tradingcalendar.py @@ -84,8 +84,7 @@ def earlier_in_day(d1, d2): """ Return true if d1 falls earlier in its own day than d2. """ - d1 = d1.replace(year = d2.year, day = d2.day) - return d1 < d2 + return d1.time() < d2.time() WEEKDAYS = [rrule.MO, rrule.TU, rrule.WE, rrule.TH, rrule.FR] @@ -97,7 +96,8 @@ market_opens_with_holidays = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart=datetime(1970, 1, 1, tzinfo = pytz.utc), + dtstart=datetime(2000, 1, 1, tzinfo = pytz.utc), + until=datetime(2014 , 1, 1, tzinfo = pytz.utc) ) # Recurrence rule that generates all market closes since Jan 1, 1970. @@ -108,7 +108,8 @@ market_closes_with_holidays = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart=datetime(1970, 1, 1, tzinfo = pytz.utc), + dtstart=datetime(2001, 1, 1, tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rules for excluding the market open/close on new years. @@ -118,7 +119,8 @@ new_years_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) new_years_closes = rrule.rrule( rrule.MONTHLY, @@ -126,7 +128,8 @@ new_years_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rules for excluding MLK day. It is always the third @@ -138,7 +141,8 @@ mlk_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) mlk_closes = rrule.rrule( rrule.MONTHLY, @@ -147,7 +151,8 @@ mlk_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rules for generating the market open/close for @@ -160,7 +165,8 @@ presidents_day_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) presidents_day_closes = rrule.rrule( rrule.MONTHLY, @@ -169,7 +175,8 @@ presidents_day_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rules for generating the market open/close for good @@ -181,7 +188,8 @@ good_friday_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) good_friday_closes = rrule.rrule( rrule.DAILY, @@ -189,7 +197,8 @@ good_friday_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rules for generating the market open/close for memorial @@ -201,7 +210,8 @@ memorial_day_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) memorial_day_closes = rrule.rrule( rrule.MONTHLY, @@ -210,7 +220,8 @@ memorial_day_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rules for generating the market open/close for July 4th. @@ -221,7 +232,8 @@ july_4th_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) july_4th_closes = rrule.rrule( rrule.MONTHLY, @@ -230,7 +242,8 @@ july_4th_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rule for generating the market open/close for labor day. @@ -242,7 +255,8 @@ labor_day_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) labor_day_closes = rrule.rrule( rrule.MONTHLY, @@ -251,7 +265,8 @@ labor_day_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence rule for generating the market open/close for @@ -264,7 +279,8 @@ thanksgiving_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) thanksgiving_closes = rrule.rrule( rrule.MONTHLY, @@ -273,7 +289,8 @@ thanksgiving_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) # Recurrence relation for generating the market open/close for @@ -286,7 +303,8 @@ christmas_opens = rrule.rrule( byhour = 14, byminute = 30, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) christmas_closes = rrule.rrule( rrule.MONTHLY, @@ -295,8 +313,10 @@ christmas_closes = rrule.rrule( byhour = 21, byminute = 0, cache = True, - dtstart = datetime(1970, 1,1,tzinfo = pytz.utc) + dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), + until=datetime(2014, 1, 1, tzinfo = pytz.utc) ) + # All NYSE observed holidays. holiday_opens = [ new_years_opens, @@ -322,12 +342,15 @@ holiday_closes = [ ] # Valid market opens are given by all market opens minus holidays. -opens = rrule.rruleset() +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() +closes = rrule.rruleset(cache=True) closes.rrule(market_closes_with_holidays) for holiday_rule in holiday_closes: closes.exrule(holiday_rule) +close_count = closes.count() +