diff --git a/zipline/gens/mavg.py b/zipline/gens/mavg.py index 4d477c9d..1978f92e 100644 --- a/zipline/gens/mavg.py +++ b/zipline/gens/mavg.py @@ -1,4 +1,3 @@ - from numbers import Number from datetime import datetime, timedelta from collections import defaultdict @@ -8,15 +7,15 @@ from zipline.gens.transform import EventWindow class MovingAverage(object): """ - Class that maintains a dictionary from sids to EventWindows - Upon receipt of each message we update the - corresponding window and return the calculated average. + Class that maintains a dictionary from sids to + MovingAverageEventWindows. For each sid, we maintain moving + averages over any number of distinct fields (For example, we can + maintain a sid's average volume as well as its average price.) """ def __init__(self, delta, fields): self.delta = delta self.fields = fields - # 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) @@ -27,13 +26,9 @@ class MovingAverage(object): def update(self, event): """ - Update the event window for this event's sid. Return an ndict from - tracked fields to averages. + Update the event window for this event's sid. Return an ndict + from tracked fields to moving averages. """ - assert isinstance(event, ndict),"Bad event in MovingAverage: %s" % event - assert event.has_key('sid'), "No sid in MovingAverage: %s" % event - assert event.has_key('dt'), "No dt in MovingAverage: %s" % event - # This will create a new EventWindow if this is the first # message for this sid. window = self.sid_windows[event.sid] @@ -42,22 +37,34 @@ class MovingAverage(object): class MovingAverageEventWindow(EventWindow): """ - Calculates a moving average over all specified fields. + Iteratively calculates moving averages for a particular sid over a + given time window. We can maintain averages for arbitrarily many + fields on a single sid. (For example, we might track average + price as well as average volume for a single sid.) The expected + functionality of this class is to be instantiated inside a + MovingAverage transform. """ - # Subclass initializer. The superclass also requires a timedelta - # argument, so instantiation should look like: - # mavg = MovingAverageEventWindow(timedelta(minutes=1), ['price']) - def init(self, fields): + + def __init__(self, delta, fields): + + # Call the superclass constructor to set up base EventWindow + # infrastructure. + EventWindow.__init__(self, delta) + + # We maintain a dictionary of totals for each of our tracked + # fields. self.fields = fields self.totals = defaultdict(float) + # Subclass customization for adding new events. def handle_add(self, event): # Sanity check on the event. - self.assert_all_fields(event) + self.assert_required_fields(event) # Increment our running totals with data from the event. for field in self.fields: self.totals[field] += event[field] + # Subclass customization for removing expired events. def handle_remove(self, event): # Decrement our running totals with data from the event. for field in self.fields: @@ -65,12 +72,12 @@ class MovingAverageEventWindow(EventWindow): def average(self, field): """ - Calculate the average value of our ticks over a given field. + Calculate the average value of our ticks over a single field. """ # Sanity check. assert field in self.fields - # Averages are 0 by convention if we have no ticks. + # Averages are None by convention if we have no ticks. if len(self.ticks) == 0: return 0.0 @@ -82,15 +89,14 @@ class MovingAverageEventWindow(EventWindow): """ Return an ndict of all our tracked averages. """ - out = ndict() - + out = ndict() for field in self.fields: out[field] = self.average(field) return out - def assert_all_fields(self, event): + def assert_required_fields(self, event): """ - We only track events with all the fields we care about. + We only allow events with all of our tracked fields. """ for field in self.fields: assert event.has_key(field), \ diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index b1a0ed96..7b86da05 100644 --- a/zipline/gens/tradegens.py +++ b/zipline/gens/tradegens.py @@ -25,9 +25,9 @@ def mock_prices(count, rand = False): """ if rand: - return (random.uniform(0.0, 10.0) for i in xrange(count)) + return (random.uniform(1.0, 10.0) for i in xrange(count)) else: - return (float(i % 11) for i in xrange(1,count+1)) + return (float(i % 10) + 1.0 for i in xrange(count)) def mock_volumes(count, rand = False): """ diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index bdee3270..a191e50b 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -156,15 +156,10 @@ class EventWindow: # Mark this as an abstract base class. __metaclass__ = ABCMeta - def __init__(self, delta, *args, **kwargs): + def __init__(self, delta): self.ticks = deque() self.delta = delta - self.init(*args, **kwargs) - @abstractmethod - def init(self): - raise NotImplementedError() - @abstractmethod def handle_add(self, event): raise NotImplementedError() @@ -193,7 +188,9 @@ class EventWindow: # Subclasses should override handle_remove to define # behavior for removing ticks. self.handle_remove(popped) - + + # All event windows expect to receive events with datetime fields + # that arrive in sorted order. def assert_well_formed(self, event): assert isinstance(event, ndict), "Bad event in EventWindow:%s" % event assert event.has_key('dt'), "Missing dt in EventWindow:%s" % event diff --git a/zipline/gens/vwap.py b/zipline/gens/vwap.py new file mode 100644 index 00000000..4e8a2f8b --- /dev/null +++ b/zipline/gens/vwap.py @@ -0,0 +1,82 @@ +from numbers import Number +from datetime import datetime, timedelta +from collections import defaultdict + +from zipline import ndict +from zipline.gens.transform import EventWindow + +class VWAP(object): + """ + Class that maintains a dictionary from sids to VWAPEventWindows. + """ + def __init__(self, delta): + self.delta = delta + + # 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) + + def update(self, event): + """ + Update the event window for this event's sid. Returns the + current vwap for the sid. + """ + # 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_vwap() + + +class VWAPEventWindow(EventWindow): + """ + Iteratively maintains a vwap for a single sid over a given + timedelta. + """ + def __init__(self, delta): + EventWindow.__init__(self, delta) + self.flux = 0.0 + self.totalvolume = 0.0 + + # Subclass customization for adding new events. + def handle_add(self, event): + # Sanity check on the event. + self.assert_required_fields(event) + self.flux += event.volume * event.price + self.totalvolume += event.volume + + # Subclass customization for removing expired events. + def handle_remove(self, event): + self.flux -= event.volume * event.price + self.totalvolume -= event.volume + + def get_vwap(self): + """ + Return the calculated vwap for this sid. + """ + # By convention, vwap is None if we have no events. + if len(self.ticks) == 0 + return None + else: + return (self.flux / self.totalvolume) + + # We need numerical price and volume to calculate a vwap. + def assert_required_fields(self, event): + assert isinstance(event.price, Number) + assert isinstance(event.volume, Number) + +if __name__ == "__main__": + + from zipline.gens.tradegens import SpecificEquityTrades + from zipline.gens.transform import StatefulTransform + + source = SpecificEquityTrades() + vwap = StatefulTransform(VWAP, timedelta(minutes = 10)) + + out = vwap.transform(source) + +