diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index d757c04a..78b1bea8 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -89,3 +89,17 @@ class TestTransformAlgorithm(TestCase): assert algo.registered_transforms['mavg']['kwargs'] == \ {'window_length': 2, 'market_aware': True} assert algo.registered_transforms['mavg']['class'] is MovingAverage + + def test_data_frequency_setting(self): + algo = TestRegisterTransformAlgorithm(data_frequency='daily') + self.assertEqual(algo.data_frequency, 'daily') + self.assertEqual(algo.annualizer, 250) + + algo = TestRegisterTransformAlgorithm(data_frequency='minute') + self.assertEqual(algo.data_frequency, 'minute') + self.assertEqual(algo.annualizer, 250 * 6 * 60) + + algo = TestRegisterTransformAlgorithm(data_frequency='minute', + annualizer=10) + self.assertEqual(algo.data_frequency, 'minute') + self.assertEqual(algo.annualizer, 10) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index dda384fb..f0ae7b8b 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -32,7 +32,7 @@ from zipline.finance.slippage import ( transact_partial ) from zipline.finance.commission import PerShare, PerTrade - +from zipline.finance.constants import ANNUALIZER from zipline.gens.composites import ( date_sorted_sources, @@ -65,8 +65,16 @@ class TradingAlgorithm(object): """ def __init__(self, *args, **kwargs): - """ - Initialize sids and other state variables. + """Initialize sids and other state variables. + + :Arguments: + data_frequency : str (daily, hourly or minutely) + The duration of the bars. + annualizer : int + Which constant to use for annualizing risk metrics. + If not provided, will extract from data_frequency. + capital_base : float + How much capital to start with. """ self.done = False self.order = None @@ -84,11 +92,20 @@ class TradingAlgorithm(object): self.slippage = VolumeShareSlippage() self.commission = PerShare() + if 'data_frequency' in kwargs: + self.set_data_frequency(kwargs.pop('data_frequency')) + else: + self.data_frequency = None + + # Override annualizer if set + if 'annualizer' in kwargs: + self.annualizer = kwargs['annualizer'] + # set the capital base self.capital_base = kwargs.get('capital_base', DEFAULT_CAPITAL_BASE) - # an algorithm subclass needs to set initialized to True - # when it is fully initialized. + # an algorithm subclass needs to set initialized to True when + # it is fully initialized. self.initialized = False # call to user-defined constructor method @@ -295,3 +312,8 @@ class TradingAlgorithm(object): def set_transforms(self, transforms): assert isinstance(transforms, list) self.transforms = transforms + + def set_data_frequency(self, data_frequency): + assert data_frequency in ('daily', 'minute') + self.data_frequency = data_frequency + self.annualizer = ANNUALIZER[self.data_frequency] diff --git a/zipline/finance/constants.py b/zipline/finance/constants.py new file mode 100644 index 00000000..ce7ec8f6 --- /dev/null +++ b/zipline/finance/constants.py @@ -0,0 +1,23 @@ +# +# Copyright 2012 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. + +TRADING_DAYS_IN_YEAR = 250 +TRADING_HOURS_IN_DAY = 6 +MINUTES_IN_HOUR = 60 + +ANNUALIZER = {'daily': TRADING_DAYS_IN_YEAR, + 'hourly': TRADING_DAYS_IN_YEAR * TRADING_HOURS_IN_DAY, + 'minute': TRADING_DAYS_IN_YEAR * TRADING_HOURS_IN_DAY * + MINUTES_IN_HOUR}