diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index b3fa7576..4a74b8f6 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -88,6 +88,29 @@ def sequential_transforms(stream_in, *transforms): dt_aliased = alias_dt(stream_out) return add_done(dt_aliased) +def sequential_transforms_dict(stream_in, transforms): + """ + Apply each transform in transforms sequentially to each event in stream_in. + Each transform application will add a new entry indexed to the transform's + hash string. + """ + + assert isinstance(transforms, dict) + + for tnfm in transforms.itervalues(): + tnfm.forward_all = False + tnfm.update_in_place = False + tnfm.append_value = True + + # Recursively apply all transforms to the stream. + stream_out = reduce(lambda stream, tnfm: tnfm.transform(stream), + transforms, + stream_in) + + dt_aliased = alias_dt(stream_out) + return add_done(dt_aliased) + + def alias_dt(stream_in): """ Alias the dt field to datetime on each message. diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index 2e8f6bea..fad8bc29 100644 --- a/zipline/gens/tradegens.py +++ b/zipline/gens/tradegens.py @@ -97,8 +97,9 @@ class SpecificEquityTrades(object): return self.__class__.__name__ + "-" + self.arg_string def create_fresh_generator(self): - if self.event_list: + for event in self.event_list: + event['source_id'] = self.get_hash() unfiltered = (event for event in self.event_list) # Set up iterators for each expected field. diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 651c337d..099fc1eb 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -93,10 +93,10 @@ class StatefulTransform(object): #TODO: refactor this to avoid unnecessary copying. assert_sort_unframe_protocol(message) - message_copy = deepcopy(message) - + #message_copy = deepcopy(message) + message_copy = message # Same shared pointer issue here as above. - tnfm_value = self.state.update(deepcopy(message_copy)) + tnfm_value = self.state.update(message_copy) # FORWARDER flag means we want to keep all original # values, plus append tnfm_id and tnfm_value. Used for diff --git a/zipline/lines.py b/zipline/lines.py index 2a59190f..c05c5c1c 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -63,6 +63,11 @@ import sys import zmq import os from signal import SIGHUP, SIGINT +import datetime +import pytz +import pandas as pd +import numpy as np + import multiprocessing from setproctitle import setproctitle @@ -70,6 +75,10 @@ from zipline.test_algorithms import TestAlgorithm from zipline.finance.trading import SIMULATION_STYLE from zipline.utils.log_utils import ZeroMQLogHandler, stdout_only_pipe from zipline.utils import factory +from zipline.utils.factory import create_trading_environment +from zipline.gens.tradegens import SpecificEquityTrades +from zipline import ndict +from zipline.protocol import DATASOURCE_TYPE from zipline.test_algorithms import TestAlgorithm @@ -358,3 +367,98 @@ class SimulatedTrading(object): #------------------- return sim + + +def create_sp_source(start_dt=None, end_dt=None): + if start_dt is None: + start_dt = datetime.datetime(2002, 1, 1, tzinfo=pytz.utc) + if end_dt is None: + end_dt = datetime.datetime(2008, 1, 1, tzinfo=pytz.utc) + + sp_events, _ = factory.load_market_data() + sp_transformed = [] + for event in sp_events: + transformed = ndict(event.to_dict()) + if (transformed.dt < start_dt) or (transformed.dt > end_dt): + continue + transformed['sid'] = 0 + transformed['price'] = transformed['returns'] + transformed['type'] = DATASOURCE_TYPE.TRADE + sp_transformed.append(transformed) + + source = SpecificEquityTrades(event_list=sp_transformed) + + return source + +class Zipline(object): + def __init__(self, **kwargs): + algorithm = kwargs.get('algorithm', TestAlgorithm) + source_descrs = kwargs.get('sources', ['S&P']) + if isinstance(source_descrs, str): + source_descrs = [source_descrs] + + sources = [] + for source_descr in source_descrs: + if isinstance(source_descr, str): + if source_descr == 'S&P': + source = create_sp_source() + else: + raise NotImplementedError, "Source with name {source_descr} not known.".format(source_descr=source_descr) + else: + source = source_descr + + sources.append(source) + + environment = kwargs.get('environment', create_trading_environment()) + + try: + transform_descrs = kwargs.get('transforms', algorithm.registered_transforms) + except: + print "Couldn't load any registered_transforms." + transform_descrs = {} + + # Create transforms by wrapping them into StatefulTransforms + transforms = [] + for namestring, trans_descr in transform_descrs.iteritems(): + sf = StatefulTransform( + trans_descr['class'], + *trans_descr['args'], + **trans_descr['kwargs'] + ) + sf.namestring = namestring + + transforms.append(sf) + + results_socket_uri = None + context = None + sim_id = None + style = SIMULATION_STYLE.FIXED_SLIPPAGE + + self.simulated_trading = SimulatedTrading( + sources, + transforms, + algorithm, + environment, + style, + results_socket_uri, + context, + sim_id) + + + def run(self): + # drain simulated_trading + perfs = [perf for perf in self.simulated_trading] + + # create daily stats dataframe + daily_perfs = [] + cum_perfs = [] + for perf in perfs: + if 'daily_perf' in perf: + daily_perfs.append(perf['daily_perf']) + else: + cum_perfs.append(perf) + + daily_dts = [np.datetime64(perf['period_close'], utc=True) for perf in daily_perfs] + daily_stats = pd.DataFrame(daily_perfs, index=daily_dts) + + return daily_stats diff --git a/zipline/optimize/algorithms.py b/zipline/optimize/algorithms.py index c4e3fb9d..119c0f76 100644 --- a/zipline/optimize/algorithms.py +++ b/zipline/optimize/algorithms.py @@ -1,3 +1,6 @@ +from zipline.gens.mavg import MovingAverage +from datetime import datetime, timedelta + class BuySellAlgorithm(object): """Algorithm that buys and sells alternatingly. The amount for each order can be specified. In addition, an offset that will @@ -46,3 +49,57 @@ 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 Algorithm(object): + def set_order(self, order_callable): + self.order = order_callable + + def get_sid_filter(self): + return [self.sid] + + def set_logger(self, logger): + self.logger = logger + + def initialize(self): + 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} + + +# Inherits from Algorithm base class +class DMA(Algorithm): + """Dual Moving Average algorithm. + """ + + def __init__(self, sid, amount, short_window=20, long_window=40): + self.sid = sid + self.amount = amount + self.done = False + self.order = None + self.frame_count = 0 + self.portfolio = None + self.orders = [] + self.market_entered = False + self.prices = [] + self.events = 0 + self.add_transform(MovingAverage, 'short_mavg', ['price'], market_aware=False, delta=timedelta(days=short_window)) + self.add_transform(MovingAverage, 'long_mavg', ['price'], market_aware=False, delta=timedelta(days=long_window)) + + def handle_data(self, data): + self.events += 1 + # access transforms via their user-defined tag + if (data[self.sid].short_mavg > data[self.sid].long_mavg) and not self.market_entered: + self.order(self.sid, 100) + self.market_entered = True + elif (data[self.sid].short_mavg < data[self.sid].long_mavg) and self.market_entered: + self.order(self.sid, -100) + self.market_entered = False diff --git a/zipline/optimize/example.py b/zipline/optimize/example.py new file mode 100644 index 00000000..0119161b --- /dev/null +++ b/zipline/optimize/example.py @@ -0,0 +1,19 @@ +from zipline.lines import Zipline +from zipline.optimize.algorithms import DMA +import pandas as pd +import matplotlib.pyplot as plt +import cProfile + +def run(): + myalgo = DMA(sid=0, amount=100) + zp = Zipline(algorithm=myalgo, sources='S&P') + stats = zp.run() + print stats + return stats + + +#cProfile.run('run()') + +stats = run() +stats.returns.plot() +plt.show() \ No newline at end of file diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index a7881fa8..723c5fb9 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -46,6 +46,22 @@ The algorithm must expose methods: """ +# 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 Algorithm(object): + def set_order(self, order_callable): + self.order = order_callable + + def get_sid_filter(self): + return [self.sid] + + def add_transform(self, transform_class, tag, **kwargs): + if not hasattr(self, 'registered_transforms'): + self.registered_transforms = {} + + self.registered_transforms[tag] = transform_class(**kwargs) + class TestAlgorithm(): """