From b74721421e5dde6dbd1f4608ed991614fe522a3e Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Fri, 7 Sep 2012 16:55:01 -0400 Subject: [PATCH] WIP: First naive implementation of BatchWindow. --- zipline/gens/transform.py | 91 +++++++++++++++++++++++++++++++++------ 1 file changed, 78 insertions(+), 13 deletions(-) diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 03eb41ea..bc8b716a 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -9,6 +9,8 @@ from datetime import datetime from collections import deque from abc import ABCMeta, abstractmethod +import pandas as pd + from zipline import ndict from zipline.utils.tradingcalendar import non_trading_days from zipline.gens.utils import assert_sort_unframe_protocol, hash_args @@ -36,7 +38,7 @@ class TransformMeta(type): still recover an instance of a "raw" Foo by introspecting the resulting StatefulTransform's 'state' field. """ - + def __call__(cls, *args, **kwargs): return StatefulTransform(cls, *args, **kwargs) @@ -53,19 +55,19 @@ class StatefulTransform(object): def __init__(self, tnfm_class, *args, **kwargs): assert isinstance(tnfm_class, (types.ObjectType, types.ClassType)), \ "Stateful transform requires a class." - assert tnfm_class.__dict__.has_key('update'), \ + assert hasattr(tnfm_class, 'update'), \ "Stateful transform requires the class to have an update method" # Flag set inside the Passthrough transform class to signify special # behavior if we are being fed to merged_transforms. - self.passthrough = tnfm_class.__dict__.get('PASSTHROUGH', False) - + self.passthrough = hasattr(tnfm_class, 'PASSTHROUGH') + # Flags specifying how to append the calculated value. # Merged is the default for ease of testing, but we use sequential # in production. self.sequential = False self.merged = True - + # Create an instance of our transform class. if isinstance(tnfm_class, TransformMeta): # Classes derived TransformMeta have their __call__ @@ -104,12 +106,12 @@ class StatefulTransform(object): continue assert_sort_unframe_protocol(message) - + # This flag is set by by merged_transforms to ensure # isolation of messages. if self.merged: message = deepcopy(message) - + tnfm_value = self.state.update(message) # PASSTHROUGH flag means we want to keep all original @@ -133,7 +135,7 @@ class StatefulTransform(object): out_message.tnfm_value = tnfm_value out_message.dt = message.dt yield out_message - + # Sequential flag should be used to add a single new # key-value pair to the event. The new key is this # transform's namestring, and its value is the value @@ -147,9 +149,11 @@ class StatefulTransform(object): out_message = message out_message[self.namestring] = tnfm_value yield out_message - + log.info('Finished StatefulTransform [%s]' % self.get_hash()) + + class EventWindow(object): """ Abstract base class for transform classes that calculate iterative @@ -218,7 +222,7 @@ class EventWindow(object): # Subclasses should override handle_add to define behavior for # adding new ticks. self.handle_add(event) - + if self.market_aware: self.add_new_holidays(event.dt) @@ -229,14 +233,14 @@ class EventWindow(object): # | | # 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 add_new_holidays(self, newest): # Add to our tracked window any untracked holidays that are # older than our newest event. (newest should always be @@ -256,7 +260,7 @@ class EventWindow(object): calendar_dates_between = (newest.date() - oldest.date()).days holidays_between = len(self.cur_holidays) trading_days_between = calendar_dates_between - holidays_between - + # "Put back" a day if oldest is earlier in its day than newest, # reflecting the fact that we haven't yet completed the last # day in the window. @@ -277,3 +281,64 @@ class EventWindow(object): # 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]) + + +class BatchWindow(EventWindow): + def __init__(self, func, refresh_period=None, wind_length=None, sids=None): + super(BatchWindow, self).__init__(True, days=wind_length, delta=None) + self.func = func + self.sids = sids + self.refresh_period = refresh_period + self.wind_length = wind_length + + self.last_calc = False + self.full = False + self.last_refresh = None + + self.updated = False + + def handle_data(self, data): + """ + New method to handle a data frame as sent to the algorithm's handle_data + method. + """ + dts = [data[sid].datetime for sid in self.sids] + prices = [data[sid].price for sid in self.sids] + volumes = [data[sid].volume for sid in self.sids] + + price_df = pd.DataFrame(prices, columns=self.sids, index=dts) + volume_df = pd.DataFrame(volumes, columns=self.sids, index=dts) + + event = ndict({ + 'dt' : max(dts), + 'prices': price_df, + 'volumes': volume_df, + }) + + self.update(event) + + def handle_add(self, event): + if not self.last_calc: + self.last_calc = event.dt + return + + age = event.dt - self.last_refresh + if age.days >= self.refresh_period: + self.prices = pd.concat(self.ticks.prices) + self.volumes = pd.concat(self.ticks.volumes) + + self.updated = True + else: + self.updated = False + + self.last_refresh = event.dt + + def handle_remove(self, event): + # since an event is expiring, we know the window is full + self.full = True + + def __call__(self, *args, **kwargs): + if self.updated: + self.cached = self.func(self.prices, self.volumes, *args, **kwargs) + + return self.cached