From 280f122353d5bef8164d44fd0ed3b26342bf5624 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Mon, 17 Sep 2012 18:35:21 -0400 Subject: [PATCH] WIP: Lot of refactoring and bugfixing. --- zipline/gens/cov.py | 12 ++- zipline/gens/tradegens.py | 12 +-- zipline/gens/tradesimulation.py | 4 +- zipline/gens/transform.py | 87 +++++++++++------- zipline/optimize/algorithms.py | 152 +++++++++++++++++++++++--------- zipline/optimize/example.py | 28 +++--- zipline/test_algorithms.py | 2 +- 7 files changed, 199 insertions(+), 98 deletions(-) diff --git a/zipline/gens/cov.py b/zipline/gens/cov.py index eefb1c96..85985795 100644 --- a/zipline/gens/cov.py +++ b/zipline/gens/cov.py @@ -1,5 +1,9 @@ -from zipline.gens.transform import EventWindowBatch +from zipline.gens.transform import BatchWindow, batch_transform -class CovEventWindow(EventWindowBatch): - def get_value(self, prices, volumes): - return prices.cov() \ No newline at end of file +class CovEventWindow(BatchWindow): + def get_value(self, data): + return data.cov() + +@batch_transform +def cov(data): + return data.cov() diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index e7c4e375..652ea135 100644 --- a/zipline/gens/tradegens.py +++ b/zipline/gens/tradegens.py @@ -181,11 +181,11 @@ class DataFrameSource(SpecificEquityTrades): self.data = data # Unpack config dictionary with default values. - self.count = kwargs.get('count', 500) - self.sids = kwargs.get('sids', [0]) - self.start = kwargs.get('start', datetime(1957, 1, 1, 0, tzinfo = pytz.utc)) - self.end = kwargs.get('end', datetime(2010, 1, 1, tzinfo=pytz.utc)) - self.delta = kwargs.get('delta', timedelta(days = 1)) + self.count = kwargs.get('count', len(data)) + self.sids = kwargs.get('sids', data.columns) + self.start = kwargs.get('start', data.index[0]) + self.end = kwargs.get('end', data.index[-1]) + self.delta = kwargs.get('delta', data.index[1]-data.index[0]) # Default to None for event_list and filter. self.filter = kwargs.get('filter') @@ -207,7 +207,7 @@ class DataFrameSource(SpecificEquityTrades): for sid, price in series.iterkv(): event = copy(event) - event['sid'] = 0 + event['sid'] = sid event['price'] = price yield ndict(event) diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index cca56e18..61527c1a 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -2,6 +2,7 @@ from logbook import Logger, Processor from datetime import datetime from itertools import groupby +from operator import attrgetter from zipline import ndict from zipline.utils.timeout import Heartbeat, Timeout @@ -226,8 +227,7 @@ class AlgorithmSimulator(object): # Group together events with the same dt field. This depends on the # events already being sorted. - for date, snapshot in groupby(stream_in, lambda e: e.dt): - + for date, snapshot in groupby(stream_in, attrgetter('dt')): # Set the simulation date to be the first event we see. # This should only occur once, at the start of the test. if self.simulation_dt == None: diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 87dad979..e4171ebc 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -217,7 +217,7 @@ class EventWindow(object): self.assert_well_formed(event) # Add new event and increment totals. - self.ticks.append(event) + self.ticks.append(deepcopy(event)) # Subclasses should override handle_add to define behavior for # adding new ticks. @@ -266,6 +266,7 @@ class EventWindow(object): # day in the window. if oldest.time() > newest.time(): trading_days_between -= 1 + return trading_days_between >= self.days def out_of_delta(self, oldest, newest): @@ -284,62 +285,84 @@ class EventWindow(object): 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) + def __init__(self, func=None, refresh_period=None, days=None, sids=None): + super(BatchWindow, self).__init__(True, days=days, delta=None) + self.func = func self.sids = sids self.refresh_period = refresh_period - self.wind_length = wind_length + self.days = days - self.last_calc = False self.full = False self.last_refresh = None self.updated = False + self.data = None - # 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] + def handle_data(self, data): + """ + New method to handle a data frame as sent to the algorithm's handle_data + method. + """ + # extract dates + dts = [data[sid].datetime for sid in self.sids] + # we have to provide the event with a dt. This is only for + # checking if the event is outside the window or not so a + # couple of seconds shouldn't matter + data.dt = max(dts) - # price_df = pd.DataFrame(prices, columns=self.sids, index=dts) - # volume_df = pd.DataFrame(volumes, columns=self.sids, index=dts) + # append data frame to window + self.update(data) - # event = ndict({ - # 'dt' : max(dts), - # 'prices': price_df, - # 'volumes': volume_df, - # }) - - # self.update(event) + # return newly computed or cached value + return self.compute() def handle_add(self, event): - import pdb; pdb.set_trace() - - if not self.last_calc: - self.last_calc = event.dt + if not self.last_refresh: + self.last_refresh = 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) + # create Series price object + data_sids = {} + for sid in self.sids: + dts = [tick[sid].dt for tick in self.ticks] + prices = [tick[sid].price for tick in self.ticks] + data_sids[sid] = pd.Series(prices, index=dts) + + # concatenate different sids into one df + self.data = pd.concat(data_sids, axis=1) self.updated = True + self.last_refresh = event.dt 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): + def get_value(self, *args, **kwargs): + raise NotImplementedError("Either overwrite get_value or provide a func argument.") + + def compute(self, *args, **kwargs): + if self.data is None: + return False + if self.updated: - self.cached = self.get_value(self.prices, self.volumes, *args, **kwargs) + if self.func is not None: + # user supplied function + self.cached = self.func(self.data, *args, **kwargs) + else: + # assume inheritance + self.cached = self.get_value(self.data, *args, **kwargs) return self.cached + + +# decorator for BatchWindow +def batch_transform(func): + def create_transform(*args, **kwargs): + return BatchWindow(*args, func=func, **kwargs) + + return create_transform diff --git a/zipline/optimize/algorithms.py b/zipline/optimize/algorithms.py index 0a860bd2..218dc989 100644 --- a/zipline/optimize/algorithms.py +++ b/zipline/optimize/algorithms.py @@ -48,6 +48,7 @@ class BuySellAlgorithm(object): self.portfolio = portfolio def handle_data(self, frame): + print frame.sid order_size = self.buy_or_sell * (self.amount - (self.offset**2)) self.order(self.sid, order_size) @@ -62,40 +63,114 @@ class BuySellAlgorithm(object): def get_sid_filter(self): return [self.sid] -# Algorithm base class, user algorithms inherit from this as they -# don't want to have to copy and know about set_order and -# set_portfolio -class TradingAlgorithm(object): - def _setup(self): - assert hasattr(self, 'source'), 'source not set.' - assert hasattr(self, 'sids'), "sids not set." - environment = create_trading_environment(start=self.data.index[0], end=self.data.index[-1]) +class TradingAlgorithm(object): + """ + Base class for trading algorithms. Inherit and overload handle_data(data). + + A new algorithm could look like this: + ``` + class MyAlgo(TradingAlgorithm): + def initialize(amount): + self.amount = amount + + def handle_data(data): + sid = self.sids[0] + self.order(sid, amount) + ``` + To then run this algorithm: + + >>> my_algo = MyAlgo(100) + >>> stats = my_algo.run(data) + + """ + def __init__(self, sids, *args, **kwargs): + """ + Initialize sids and other state variables. + + Calls user-defined initialize and forwarding *args and **kwargs. + """ + self.sids = sids + self.done = False + self.order = None + self.frame_count = 0 + self.portfolio = None + + self.registered_transforms = {} + + # call to user-defined initialize method + self.initialize(*args, **kwargs) + + def _create_simulator(self, source): + """ + Create trading environment, transforms and SimulatedTrading object. + + Gets called by self.run(data). + """ + environment = create_trading_environment(start=source.data.index[0], end=source.data.index[-1]) # Create transforms by wrapping them into StatefulTransforms transforms = [] - if hasattr(self, 'registered_transforms'): - for namestring, trans_descr in self.registered_transforms.iteritems(): - sf = StatefulTransform( - trans_descr['class'], - *trans_descr['args'], - **trans_descr['kwargs'] - ) - sf.namestring = namestring + for namestring, trans_descr in self.registered_transforms.iteritems(): + sf = StatefulTransform( + trans_descr['class'], + *trans_descr['args'], + **trans_descr['kwargs'] + ) + sf.namestring = namestring - transforms.append(sf) + transforms.append(sf) - - self.simulated_trading = SimulatedTrading( - [self.source], + # SimulatedTrading is the main class handling data streaming, + # application of transforms and calling of the user algo. + return SimulatedTrading( + [source], transforms, self, environment, FixedSlippage() ) + def run(self, data): + """ + Run the algorithm. + + :Arguments: + data : pandas.DataFrame + * columns must consist of ints representing the different sids + * index must be TimeStamps + * array contents should be price + + :Returns: + daily_stats : pandas.DataFrame + Daily performance metrics such as returns, alpha etc. + + """ + assert isinstance(data, pd.DataFrame) + assert isinstance(data.index, pd.Timeseries) + + source = DataFrameSource(data, sids=self.sids) + + # create transforms and zipline + simulated_trading = self._create_simulator(source) + + # loop through simulated_trading, each iteration returns a + # perf ndict + perfs = [] + for perf in simulated_trading: + #from nose.tools import set_trace; set_trace() + perfs.append(perf) + + #perfs = list(self.simulated_trading) + + # convert perf ndict to pandas dataframe + daily_stats = self._create_daily_stats(perfs) + + return daily_stats + + def _create_daily_stats(self, perfs): - # create daily stats dataframe + # create daily and cumulative stats dataframe daily_perfs = [] cum_perfs = [] for perf in perfs: @@ -109,21 +184,23 @@ class TradingAlgorithm(object): return daily_stats - def run(self, data, compute_risk_metrics=False): - self.source = DataFrameSource(data, sids=self.sids) - self.data = data - self._setup() + def add_transform(self, transform_class, tag, *args, **kwargs): + """Add a single-sid, sequential transform to the model. - # drain simulated_trading - perfs = [] - for perf in self.simulated_trading: - #from nose.tools import set_trace; set_trace() - perfs.append(perf) + :Arguments: + transform_class : class + Which transform to use. E.g. mavg. + tag : str + How to name the transform. Can later be access via: + data[sid].tag() - #perfs = list(self.simulated_trading) + Extra args and kwargs will be forwarded to the transform + instantiation. - daily_stats = self._create_daily_stats(perfs) - return daily_stats + """ + self.registered_transforms[tag] = {'class': transform_class, + 'args': args, + 'kwargs': kwargs} def set_portfolio(self, portfolio): self.portfolio = portfolio @@ -137,19 +214,12 @@ class TradingAlgorithm(object): def set_logger(self, logger): self.logger = logger - def initialize(self): + def initialize(self, *args, **kwargs): pass def set_slippage_override(self, slippage_callable): pass - def add_transform(self, transform_class, tag, *args, **kwargs): - if not hasattr(self, 'registered_transforms'): - self.registered_transforms = {} - - self.registered_transforms[tag] = {'class': transform_class, - 'args': args, - 'kwargs': kwargs} class BuySellAlgorithmNew(TradingAlgorithm): diff --git a/zipline/optimize/example.py b/zipline/optimize/example.py index feca6306..c25dafc9 100644 --- a/zipline/optimize/example.py +++ b/zipline/optimize/example.py @@ -6,6 +6,7 @@ import numpy as np import matplotlib.pyplot as plt import cProfile from zipline.gens.mavg import MovingAverage +from zipline.gens.cov import CovEventWindow, cov from zipline.optimize.algorithms import TradingAlgorithm from datetime import timedelta @@ -15,15 +16,9 @@ from datetime import timedelta class DMA(TradingAlgorithm): """Dual Moving Average algorithm. """ - def __init__(self, sids, amount=100, short_window=20, long_window=40): - self.sids = sids - self.amount = amount - self.done = False - self.order = None - self.frame_count = 0 - self.portfolio = None + def initialize(self, amount=100, short_window=20, long_window=40): self.orders = [] - + self.amount = amount self.prices = [] self.events = 0 @@ -33,15 +28,22 @@ class DMA(TradingAlgorithm): self.add_transform(MovingAverage, 'short_mavg', ['price'], market_aware=True, - days=short_window) #timedelta(days=int(short_window))) + days=short_window) self.add_transform(MovingAverage, 'long_mavg', ['price'], market_aware=True, - days=long_window) #timedelta(days=int(long_window))) + days=long_window) + + self.cov = CovEventWindow(sids=self.sids, refresh_period=1, days=5) + self.cov2 = cov(sids=self.sids, refresh_period=1, days=5) def handle_data(self, data): self.events += 1 + cov = self.cov.handle_data(data) + cov = self.cov2.handle_data(data) + print cov + for sid in self.sids: # access transforms via their user-defined tag if (data[sid].short_mavg['price'] > data[sid].long_mavg['price']) and not self.invested[sid]: @@ -86,8 +88,8 @@ def load_close_px(indexes=None, stocks=None): def run((short_window, long_window)): #data = pd.DataFrame.from_csv('SP500.csv') - data = load_close_px() - myalgo = DMA([0], amount=100, short_window=short_window, long_window=long_window) + data = pd.DataFrame.from_csv('aapl.csv') #load_close_px() + myalgo = DMA([0, 1], amount=100, short_window=short_window, long_window=long_window) stats = myalgo.run(data) stats['sw'] = short_window stats['lw'] = long_window @@ -153,3 +155,5 @@ def plot_returns(port_returns, bmk_returns): cum_bmk.plot(label='Benchmark') plt.title('Portfolio performance') plt.legend(loc='best') + +print run((10, 20)) \ No newline at end of file diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 3b04b243..055c8c32 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -98,7 +98,7 @@ class TestAlgorithm(): def set_slippage_override(self, slippage_callable): pass - # + class HeavyBuyAlgorithm(): """ This algorithm will send a specified number of orders, to allow unit tests