diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 1792c97a..b603b1d5 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -16,9 +16,12 @@ from unittest import TestCase from datetime import timedelta import numpy as np +from mock import MagicMock from zipline.utils.test_utils import setup_logger import zipline.utils.factory as factory +import zipline.utils.simfactory as simfactory + from zipline.test_algorithms import (TestRegisterTransformAlgorithm, RecordAlgorithm, TestOrderAlgorithm, @@ -28,13 +31,27 @@ from zipline.test_algorithms import (TestRegisterTransformAlgorithm, TestOrderPercentAlgorithm, TestTargetPercentAlgorithm, TestTargetValueAlgorithm, - EmptyPositionsAlgorithm) + EmptyPositionsAlgorithm, + initialize_noop, + handle_data_noop, + initialize_api, + handle_data_api, + noop_algo, + api_algo, + call_all_order_methods, + record_variables, + record_float_magic + ) + +from zipline.utils.test_utils import drain_zipline, assert_single_position from zipline.sources import (SpecificEquityTrades, DataFrameSource, DataPanelSource) from zipline.transforms import MovingAverage from zipline.finance.trading import SimulationParameters +from zipline.utils.api_support import set_algo_instance +from zipline.algorithm import TradingAlgorithm class TestRecordAlgorithm(TestCase): @@ -235,3 +252,238 @@ class TestPositions(TestCase): for i, expected in enumerate(expected_position_count): self.assertEqual(daily_stats.ix[i]['num_positions'], expected) + + +class TestAlgoScript(TestCase): + def setUp(self): + days = 251 + self.sim_params = factory.create_simulation_parameters(num_days=days) + setup_logger(self) + + trade_history = factory.create_trade_history( + 133, + [10.0] * days, + [100] * days, + timedelta(days=1), + self.sim_params + ) + + self.source = SpecificEquityTrades(event_list=trade_history) + self.df_source, self.df = \ + factory.create_test_df_source(self.sim_params) + + self.zipline_test_config = { + 'sid': 0, + } + + def test_noop(self): + algo = TradingAlgorithm(initialize=initialize_noop, + handle_data=handle_data_noop) + algo.run(self.df) + + def test_noop_string(self): + algo = TradingAlgorithm(script=noop_algo) + algo.run(self.df) + + def test_api_calls(self): + algo = TradingAlgorithm(initialize=initialize_api, + handle_data=handle_data_api) + algo.run(self.df) + + def test_api_calls_string(self): + algo = TradingAlgorithm(script=api_algo) + algo.run(self.df) + + def test_fixed_slippage(self): + # verify order -> transaction -> portfolio position. + # -------------- + test_algo = TradingAlgorithm( + script=""" +from zipline.api import (slippage, + commission, + set_slippage, + set_commission, + order, + record) + +def initialize(context): + model = slippage.FixedSlippage(spread=0.10) + set_slippage(model) + set_commission(commission.PerTrade(100.00)) + context.count = 1 + context.incr = 0 + +def handle_data(context, data): + if context.incr < context.count: + order(0, -1000) + record(price=data[0].price) + + context.incr += 1""", + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + # this matches the value in the algotext initialize + # method, and will be used inside assert_single_position + # to confirm we have as many transactions as orders we + # placed. + self.zipline_test_config['order_count'] = 1 + + # self.zipline_test_config['transforms'] = \ + # test_algo.transform_visitor.transforms.values() + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + + output, _ = assert_single_position(self, zipline) + + # confirm the slippage and commission on a sample + # transaction + recorded_price = output[1]['daily_perf']['recorded_vars']['price'] + transaction = output[1]['daily_perf']['transactions'][0] + self.assertEqual(100.0, transaction['commission']) + expected_spread = 0.05 + expected_commish = 0.10 + expected_price = recorded_price - expected_spread - expected_commish + self.assertEqual(expected_price, transaction['price']) + + def test_volshare_slippage(self): + # verify order -> transaction -> portfolio position. + # -------------- + test_algo = TradingAlgorithm( + script=""" +from zipline.api import * + +def initialize(context): + model = slippage.VolumeShareSlippage( + volume_limit=.3, + price_impact=0.05 + ) + set_slippage(model) + set_commission(commission.PerShare(0.02)) + context.count = 2 + context.incr = 0 + +def handle_data(context, data): + if context.incr < context.count: + # order small lots to be sure the + # order will fill in a single transaction + order(0, 5000) + record(price=data[0].price) + record(volume=data[0].volume) + record(incr=context.incr) + context.incr += 1 + """, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 100 + + # 67 will be used inside assert_single_position + # to confirm we have as many transactions as expected. + # The algo places 2 trades of 5000 shares each. The trade + # events have volume ranging from 100 to 950. The volume cap + # of 0.3 limits the trade volume to a range of 30 - 316 shares. + # The spreadsheet linked below calculates the total position + # size over each bar, and predicts 67 txns will be required + # to fill the two orders. The number of bars and transactions + # differ because some bars result in multiple txns. See + # spreadsheet for details: +# https://www.dropbox.com/s/ulrk2qt0nrtrigb/Volume%20Share%20Worksheet.xlsx + self.zipline_test_config['expected_transactions'] = 67 + + # self.zipline_test_config['transforms'] = \ + # test_algo.transform_visitor.transforms.values() + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + output, _ = assert_single_position(self, zipline) + + # confirm the slippage and commission on a sample + # transaction + per_share_commish = 0.02 + perf = output[1] + transaction = perf['daily_perf']['transactions'][0] + commish = transaction['amount'] * per_share_commish + self.assertEqual(commish, transaction['commission']) + self.assertEqual(2.029, transaction['price']) + + def test_algo_record_vars(self): + test_algo = TradingAlgorithm( + script=record_variables, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + output, _ = drain_zipline(self, zipline) + self.assertEqual(len(output), 252) + incr = [] + for o in output[:200]: + incr.append(o['daily_perf']['recorded_vars']['incr']) + + np.testing.assert_array_equal(incr, range(1, 201)) + + def test_algo_record_allow_mock(self): + """ + Test that values from "MagicMock"ed methods can be passed to record. + + Relevant for our basic/validation and methods like history, which + will end up returning a MagicMock instead of a DataFrame. + """ + test_algo = TradingAlgorithm( + script=record_variables, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + test_algo.record(foo=MagicMock()) + + def _algo_record_float_magic_should_pass(self, var_type): + test_algo = TradingAlgorithm( + script=record_float_magic % var_type, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + output, _ = drain_zipline(self, zipline) + self.assertEqual(len(output), 252) + incr = [] + for o in output[:200]: + incr.append(o['daily_perf']['recorded_vars']['data']) + np.testing.assert_array_equal(incr, [np.nan] * 200) + + def test_algo_record_nan(self): + self._algo_record_float_magic_should_pass('nan') + + def test_order_methods(self): + """Only test that order methods can be called without error. + Correct filling of orders is tested in zipline. + """ + test_algo = TradingAlgorithm( + script=call_all_order_methods, + sim_params=self.sim_params, + ) + set_algo_instance(test_algo) + + self.zipline_test_config['algorithm'] = test_algo + self.zipline_test_config['trade_count'] = 200 + + zipline = simfactory.create_test_zipline( + **self.zipline_test_config) + + output, _ = drain_zipline(self, zipline) diff --git a/tests/test_batchtransform.py b/tests/test_batchtransform.py index 8f5bdbbe..ece1c2c8 100644 --- a/tests/test_batchtransform.py +++ b/tests/test_batchtransform.py @@ -27,9 +27,10 @@ from zipline.utils.test_utils import setup_logger from zipline.sources.data_source import DataSource import zipline.utils.factory as factory +from zipline.transforms import batch_transform + from zipline.test_algorithms import (BatchTransformAlgorithm, BatchTransformAlgorithmMinute, - batch_transform, ReturnPriceBatchTransform) from zipline.algorithm import TradingAlgorithm diff --git a/zipline/__init__.py b/zipline/__init__.py index 222e4ce3..250d4537 100644 --- a/zipline/__init__.py +++ b/zipline/__init__.py @@ -26,13 +26,17 @@ from . import data from . import finance from . import gens from . import utils +from . import transforms from . algorithm import TradingAlgorithm +from . import api __all__ = [ 'data', 'finance', 'gens', 'utils', + 'transforms', + 'api', 'TradingAlgorithm' ] diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 8d5baf77..01b1fc54 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -22,7 +22,7 @@ from datetime import datetime from itertools import groupby from six.moves import filter -from six import iteritems +from six import iteritems, exec_ from operator import attrgetter from zipline.errors import ( @@ -34,6 +34,7 @@ from zipline.errors import ( from zipline.finance.performance import PerformanceTracker from zipline.sources import DataFrameSource, DataPanelSource from zipline.utils.factory import create_simulation_parameters +from zipline.utils.api_support import set_algo_instance, api_method from zipline.transforms.utils import StatefulTransform from zipline.finance.slippage import ( VolumeShareSlippage, @@ -64,19 +65,21 @@ class TradingAlgorithm(object): A new algorithm could look like this: ``` - class MyAlgo(TradingAlgorithm): - def initialize(self, sids, amount): - self.sids = sids - self.amount = amount + from zipline.api import order - def handle_data(self, data): - sid = self.sids[0] - amount = self.amount - self.order(sid, amount) + def initialize(context): + context.sid = 'AAPL' + context.amount = 100 + + def handle_data(self, data): + sid = context.sid + amount = context.amount + order(sid, amount) ``` - To then to run this algorithm: + To then to run this algorithm pass these functions to + TradingAlgorithm: - my_algo = MyAlgo([0], 100) # first argument has to be list of sids + my_algo = TradingAlgorithm(initialize, handle_data) stats = my_algo.run(data) """ @@ -84,6 +87,16 @@ class TradingAlgorithm(object): """Initialize sids and other state variables. :Arguments: + :Optional: + initialize : function + Function that is called with a single + argument at the begninning of the simulation. + handle_data : function + Function that is called with 2 arguments + (context and data) on every bar. + script : str + Algoscript that contains initialize and + handle_data function definition. data_frequency : str (daily, hourly or minutely) The duration of the bars. annualizer : int @@ -137,13 +150,46 @@ class TradingAlgorithm(object): self.portfolio_needs_update = True self._portfolio = None + # If string is passed in, execute and get reference to + # functions. + self.algoscript = kwargs.pop('script', None) + + if self.algoscript is not None: + self.ns = {} + exec_(self.algoscript, self.ns) + if 'initialize' not in self.ns: + raise ValueError('You must define an initialze function.') + if 'handle_data' not in self.ns: + raise ValueError('You must define a handle_data function.') + self._initialize = self.ns['initialize'] + self._handle_data = self.ns['handle_data'] + + # If two functions are passed in assume initialize and + # handle_data are passed in. + elif kwargs.get('initialize', False) and kwargs.get('handle_data'): + if self.algoscript is not None: + raise ValueError('You can not set script and \ + initialize/handle_data.') + self._initialize = kwargs.pop('initialize') + self._handle_data = kwargs.pop('handle_data') + # an algorithm subclass needs to set initialized to True when # it is fully initialized. self.initialized = False - # call to user-defined constructor method self.initialize(*args, **kwargs) + def initialize(self, *args, **kwargs): + # store algo reference in global space + set_algo_instance(self) + try: + self._initialize(self) + finally: + set_algo_instance(None) + + def handle_data(self, data): + self._handle_data(self, data) + def __repr__(self): """ N.B. this does not yet represent a string that can be used @@ -244,9 +290,6 @@ class TradingAlgorithm(object): """ return self._create_generator(self.sim_params) - def initialize(self, *args, **kwargs): - pass - # TODO: make a new subclass, e.g. BatchAlgorithm, and move # the run method to the subclass, and refactor to put the # generator creation logic into get_generator. @@ -324,14 +367,21 @@ class TradingAlgorithm(object): # create transforms and zipline self.gen = self._create_generator(sim_params) - # loop through simulated_trading, each iteration returns a - # perf dictionary - perfs = [] - for perf in self.gen: - perfs.append(perf) + # store algo reference in global space + set_algo_instance(self) - # convert perf dict to pandas dataframe - daily_stats = self._create_daily_stats(perfs) + try: + # loop through simulated_trading, each iteration returns a + # perf dictionary + perfs = [] + for perf in self.gen: + perfs.append(perf) + + # convert perf dict to pandas dataframe + daily_stats = self._create_daily_stats(perfs) + finally: + # remove algo from global space + set_algo_instance(None) return daily_stats @@ -375,6 +425,7 @@ class TradingAlgorithm(object): 'args': args, 'kwargs': kwargs} + @api_method def record(self, **kwargs): """ Track and record local variable (i.e. attributes) each day. @@ -382,9 +433,11 @@ class TradingAlgorithm(object): for name, value in kwargs.items(): self._recorded_vars[name] = value + @api_method def order(self, sid, amount, limit_price=None, stop_price=None): return self.blotter.order(sid, amount, limit_price, stop_price) + @api_method def order_value(self, sid, value, limit_price=None, stop_price=None): """ Place an order by desired value rather than desired number of shares. @@ -438,6 +491,7 @@ class TradingAlgorithm(object): "Algorithm expects a utc datetime" self.datetime = dt + @api_method def get_datetime(self): """ Returns a copy of the datetime. @@ -454,6 +508,7 @@ class TradingAlgorithm(object): """ self.blotter.transact = transact + @api_method def set_slippage(self, slippage): if not isinstance(slippage, SlippageModel): raise UnsupportedSlippageModel() @@ -461,6 +516,7 @@ class TradingAlgorithm(object): raise OverrideSlippagePostInit() self.slippage = slippage + @api_method def set_commission(self, commission): if not isinstance(commission, (PerShare, PerTrade, PerDollar)): raise UnsupportedCommissionModel() @@ -482,6 +538,7 @@ class TradingAlgorithm(object): self.data_frequency = data_frequency self.annualizer = ANNUALIZER[self.data_frequency] + @api_method def order_percent(self, sid, percent, limit_price=None, stop_price=None): """ Place an order in the specified security corresponding to the given @@ -492,6 +549,7 @@ class TradingAlgorithm(object): value = self.portfolio.portfolio_value * percent return self.order_value(sid, value, limit_price, stop_price) + @api_method def order_target(self, sid, target, limit_price=None, stop_price=None): """ Place an order to adjust a position to a target number of shares. If @@ -507,6 +565,7 @@ class TradingAlgorithm(object): else: return self.order(sid, target, limit_price, stop_price) + @api_method def order_target_value(self, sid, target, limit_price=None, stop_price=None): """ @@ -525,6 +584,7 @@ class TradingAlgorithm(object): else: return self.order_value(sid, target, limit_price, stop_price) + @api_method def order_target_percent(self, sid, target, limit_price=None, stop_price=None): """ @@ -547,6 +607,7 @@ class TradingAlgorithm(object): req_value = target_value - current_value return self.order_value(sid, req_value, limit_price, stop_price) + @api_method def get_open_orders(self, sid=None): if sid is None: return {key: [order.to_api_obj() for order in orders] @@ -557,10 +618,12 @@ class TradingAlgorithm(object): return [order.to_api_obj() for order in orders] return [] + @api_method def get_order(self, order_id): if order_id in self.blotter.orders: return self.blotter.orders[order_id].to_api_obj() + @api_method def cancel_order(self, order_param): order_id = order_param if isinstance(order_param, zipline.protocol.Order): diff --git a/zipline/api.py b/zipline/api.py new file mode 100644 index 00000000..bc012ab0 --- /dev/null +++ b/zipline/api.py @@ -0,0 +1,39 @@ +# +# Copyright 2013 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Note that part of the API is implemented in TradingAlgorithm as +# methods (e.g. order). These are added to this namespace via the +# decorator `api_methods` inside of algorithm.py. + +import zipline +from .finance import (commission, slippage) +from .utils import math_utils + +from zipline.finance.slippage import ( + FixedSlippage, + VolumeShareSlippage, +) + + +batch_transform = zipline.transforms.BatchTransform + +__all__ = [ + 'slippage', + 'commission', + 'math_utils', + 'batch_transform', + 'FixedSlippage', + 'VolumeShareSlippage' +] diff --git a/zipline/examples/buyapple.py b/zipline/examples/buyapple.py index 5f0a931c..e4ba859d 100755 --- a/zipline/examples/buyapple.py +++ b/zipline/examples/buyapple.py @@ -26,6 +26,9 @@ class BuyApple(TradingAlgorithm): # inherit from TradingAlgorithm """This is the simplest possible algorithm that does nothing but buy 1 apple share on each event. """ + def initialize(self): + pass + def handle_data(self, data): # overload handle_data() method self.order('AAPL', 1) # order SID (=0) and amount (=1 shares) diff --git a/zipline/examples/quantopian_buy_apple.py b/zipline/examples/quantopian_buy_apple.py new file mode 100644 index 00000000..bf0d69a2 --- /dev/null +++ b/zipline/examples/quantopian_buy_apple.py @@ -0,0 +1,45 @@ +# +# Copyright 2013 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from datetime import datetime +import pytz + +from zipline import TradingAlgorithm +from zipline.utils.factory import load_from_yahoo + +from zipline.api import order + + +def initialize(context): + context.test = 10 + + +def handle_date(context, data): + order('AAPL', 10) + print(context.test) + + +if __name__ == '__main__': + import pylab as pl + start = datetime(2008, 1, 1, 0, 0, 0, 0, pytz.utc) + end = datetime(2010, 1, 1, 0, 0, 0, 0, pytz.utc) + data = load_from_yahoo(stocks=['AAPL'], indexes={}, start=start, + end=end) + data = data.dropna() + algo = TradingAlgorithm(initialize=initialize, + handle_data=handle_date) + results = algo.run(data) + results.portfolio_value.plot() + pl.show() diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index f9a958e2..614476ea 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -78,7 +78,7 @@ from six.moves import range from six import itervalues from zipline.algorithm import TradingAlgorithm -from zipline.finance.slippage import FixedSlippage +from zipline.api import FixedSlippage class TestAlgorithm(TradingAlgorithm): @@ -665,3 +665,120 @@ class EmptyPositionsAlgorithm(TradingAlgorithm): # Should be 0 when all positions are exited. self.record(num_positions=len(self.portfolio.positions)) + + +############################## +# Quantopian style algorithms +from zipline.api import (order, + set_slippage, + record) + + +# Noop algo +def initialize_noop(context): + pass + + +def handle_data_noop(context, data): + pass + + +# API functions +def initialize_api(context): + context.incr = 0 + context.sale_price = None + set_slippage(FixedSlippage()) + + +def handle_data_api(context, data): + if context.incr == 0: + assert 0 not in context.portfolio.positions + else: + assert context.portfolio.positions[0]['amount'] == \ + context.incr, "Orders not filled immediately." + assert context.portfolio.positions[0]['last_sale_price'] == \ + data[0].price, "Orders not filled at current price." + context.incr += 1 + order(0, 1) + + record(incr=context.incr) + +########################### +# AlgoScripts as strings +noop_algo = """ +# Noop algo +def initialize(context): + pass + +def handle_data(context, data): + pass +""" + +api_algo = """ +from zipline.api import (order, + set_slippage, + FixedSlippage, + record) + +def initialize(context): + context.incr = 0 + context.sale_price = None + set_slippage(FixedSlippage()) + +def handle_data(context, data): + if context.incr == 0: + assert 0 not in context.portfolio.positions + else: + assert context.portfolio.positions[0]['amount'] == \ + context.incr, "Orders not filled immediately." + assert context.portfolio.positions[0]['last_sale_price'] == \ + data[0].price, "Orders not filled at current price." + context.incr += 1 + order(0, 1) + + record(incr=context.incr) +""" + +call_all_order_methods = """ +from zipline.api import (order, + order_value, + order_percent, + order_target, + order_target_value, + order_target_percent) + +def initialize(context): + pass + +def handle_data(context, data): + order(0, 10) + order_value(0, 300) + order_percent(0, .1) + order_target(0, 100) + order_target_value(0, 100) + order_target_percent(0, .2) +""" + +record_variables = """ +from zipline.api import record + +def initialize(context): + context.stocks = [0, 1] + context.incr = 0 + +def handle_data(context, data): + context.incr += 1 + record(incr=context.incr) +""" + +record_float_magic = """ +from zipline.api import record + +def initialize(context): + context.stocks = [0, 1] + context.incr = 0 + +def handle_data(context, data): + context.incr += 1 + record(data=float('%s')) +""" diff --git a/zipline/utils/api_support.py b/zipline/utils/api_support.py new file mode 100644 index 00000000..0a19ca1c --- /dev/null +++ b/zipline/utils/api_support.py @@ -0,0 +1,42 @@ +# +# Copyright 2013 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from functools import wraps +import zipline.api + +import threading +context = threading.local() + + +def get_algo_instance(): + return getattr(context, 'algorithm', None) + + +def set_algo_instance(algo): + context.algorithm = algo + + +def api_method(f): + # Decorator that adds the decorated class method as a callable + # function (wrapped) to zipline.api + @wraps(f) + def wrapped(*args, **kwargs): + # Get the instance and call the method + return getattr(get_algo_instance(), f.__name__)(*args, **kwargs) + # Add functor to zipline.api + setattr(zipline.api, f.__name__, wrapped) + zipline.api.__all__.append(f.__name__) + + return f