diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index cf76f715..0ad92805 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -16,6 +16,7 @@ from unittest import TestCase from datetime import timedelta import numpy as np +import pandas as pd from mock import MagicMock from zipline.utils.test_utils import setup_logger @@ -48,7 +49,9 @@ from zipline.utils.test_utils import drain_zipline, assert_single_position from zipline.sources import (SpecificEquityTrades, DataFrameSource, - DataPanelSource) + DataPanelSource, + RandomWalkSource) + from zipline.transforms import MovingAverage from zipline.finance.trading import SimulationParameters from zipline.utils.api_support import set_algo_instance @@ -214,6 +217,17 @@ class TestTransformAlgorithm(TestCase): algo.run(self.df) + def test_minute_data(self): + source = RandomWalkSource(freq='minute', + start=pd.Timestamp('2000-1-1', + tz='UTC'), + end=pd.Timestamp('2000-1-1', + tz='UTC')) + algo = TestOrderInstantAlgorithm(sim_params=self.sim_params, + data_frequency='minute', + instant_fill=True) + algo.run(source) + class TestPositions(TestCase): diff --git a/tests/test_sources.py b/tests/test_sources.py index 1ba2e937..a99e2db0 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -15,13 +15,17 @@ import pandas as pd import pytz from itertools import cycle +import numpy as np from six import integer_types from unittest import TestCase import zipline.utils.factory as factory -from zipline.sources import DataFrameSource, DataPanelSource +from zipline.sources import (DataFrameSource, + DataPanelSource, + RandomWalkSource) +from zipline.utils import tradingcalendar as calendar_nyse class TestDataFrameSource(TestCase): @@ -75,3 +79,55 @@ class TestDataFrameSource(TestCase): self.assertIn(check_field, event) self.assertTrue(isinstance(event['volume'], (integer_types))) self.assertEqual(next(stocks_iter), event['sid']) + + +class TestRandomWalkSource(TestCase): + def test_minute(self): + np.random.seed(123) + start_prices = {0: 100, + 1: 500} + start = pd.Timestamp('1990-01-01', tz='UTC') + end = pd.Timestamp('1991-01-01', tz='UTC') + source = RandomWalkSource(start_prices=start_prices, + calendar=calendar_nyse, start=start, + end=end) + self.assertIsInstance(source.start, pd.lib.Timestamp) + self.assertIsInstance(source.end, pd.lib.Timestamp) + + for event in source: + self.assertIn(event.sid, start_prices.keys()) + self.assertIn(event.dt.replace(minute=0, hour=0), + calendar_nyse.trading_days) + self.assertGreater(event.dt, start) + self.assertLess(event.dt, end) + self.assertGreater(event.price, 0, + "price should never go negative.") + self.assertEqual(event.volume, 1000) + self.assertTrue(13 <= event.dt.hour <= 21, + "event.dt.hour == %i, not during market \ + hours." % event.dt.hour) + + def test_day(self): + np.random.seed(123) + start_prices = {0: 100, + 1: 500} + start = pd.Timestamp('1990-01-01', tz='UTC') + end = pd.Timestamp('1992-01-01', tz='UTC') + source = RandomWalkSource(start_prices=start_prices, + calendar=calendar_nyse, start=start, + end=end, freq='day') + self.assertIsInstance(source.start, pd.lib.Timestamp) + self.assertIsInstance(source.end, pd.lib.Timestamp) + + for event in source: + self.assertIn(event.sid, start_prices.keys()) + self.assertIn(event.dt.replace(minute=0, hour=0), + calendar_nyse.trading_days) + self.assertGreater(event.dt, start) + self.assertLess(event.dt, end) + self.assertGreater(event.price, 0, + "price should never go negative.") + self.assertEqual(event.volume, 1000) + self.assertTrue(13 <= event.dt.hour <= 21, + "event.dt.hour == %i, not during market \ + hours." % event.dt.hour) diff --git a/zipline/sources/__init__.py b/zipline/sources/__init__.py index 2d853b4a..e88807cb 100644 --- a/zipline/sources/__init__.py +++ b/zipline/sources/__init__.py @@ -1,8 +1,9 @@ from zipline.sources.data_frame_source import DataFrameSource, DataPanelSource from zipline.sources.test_source import SpecificEquityTrades - +from .simulated import RandomWalkSource __all__ = [ 'DataFrameSource', 'DataPanelSource', - 'SpecificEquityTrades' + 'SpecificEquityTrades', + 'RandomWalkSource' ] diff --git a/zipline/sources/simulated.py b/zipline/sources/simulated.py new file mode 100644 index 00000000..aa515260 --- /dev/null +++ b/zipline/sources/simulated.py @@ -0,0 +1,138 @@ +# +# Copyright 2014 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 copy import copy +import six + +import numpy as np +from datetime import timedelta + +from zipline.sources.data_source import DataSource +from zipline.utils import tradingcalendar as calendar_nyse +from zipline.gens.utils import hash_args + + +class RandomWalkSource(DataSource): + """RandomWalkSource that emits events with prices that follow a + random walk. Will generate valid datetimes that match market hours + of the supplied calendar and can generate emit events with + user-defined frequencies (e.g. minutely). + + """ + def __init__(self, start_prices=None, freq='minute', start=None, + end=None, calendar=calendar_nyse): + """ + :Arguments: + start_prices : dict + sid -> starting price. + Default: {0: 100, 1: 500} + freq : str + Emits events according to freq. + Can be 'day' or 'minute' + start : datetime + Start dt to emit events. + end : datetime + End dt until to which emit events. + calendar : calendar object + Calendar to use. + See zipline.utils for different choices. + + :Example: + # Assumes you have instantiated your Algorithm + # as myalgo. + myalgo = MyAlgo() + source = RandomWalkSource() + myalgo.run(source) + + """ + # Hash_value for downstream sorting. + self.arg_string = hash_args(start_prices, freq, start, end, + calendar.__name__) + + self.freq = freq + if start_prices is None: + self.start_prices = {0: 100, + 1: 500} + else: + self.start_prices = start_prices + + self.calendar = calendar + if start is None: + self.start = calendar.start + else: + self.start = start + if end is None: + self.end = calendar.end_base + else: + self.end = end + + self.drift = .1 + self.sd = .1 + + self.open_and_closes = \ + calendar.open_and_closes[self.start:self.end] + + self._raw_data = None + + @property + def instance_hash(self): + return self.arg_string + + @property + def mapping(self): + return { + 'dt': (lambda x: x, 'dt'), + 'sid': (lambda x: x, 'sid'), + 'price': (float, 'price'), + 'volume': (int, 'volume'), + } + + def _gen_next_step(self, x): + x += np.random.randn() * self.sd + self.drift + return max(x, 0.1) + + def _gen_events(self, cur_prices, current_dt): + for sid, price in six.iteritems(cur_prices): + cur_prices[sid] = self._gen_next_step(cur_prices[sid]) + + event = { + 'dt': current_dt, + 'sid': sid, + 'price': cur_prices[sid], + 'volume': 1000, + } + + yield event + + def raw_data_gen(self): + cur_prices = copy(self.start_prices) + for _, (open_dt, close_dt) in self.open_and_closes.iterrows(): + current_dt = copy(open_dt) + if self.freq == 'minute': + # Emit minutely trade signals from open to close + while current_dt < close_dt: + for event in self._gen_events(cur_prices, current_dt): + yield event + current_dt += timedelta(minutes=1) + elif self.freq == 'day': + # Emit one signal per day at close + for event in self._gen_events(cur_prices, close_dt): + yield event + + @property + def raw_data(self): + if not self._raw_data: + self._raw_data = self.raw_data_gen() + return self._raw_data