From 775c564ea143aa9c43012f01816b949cc57bbe8c Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Thu, 13 Dec 2012 12:55:05 -0500 Subject: [PATCH] ENH: Added DataPanelSource. --- tests/test_algorithm.py | 12 ++++- tests/test_sources.py | 22 ++++++++- zipline/algorithm.py | 5 +- zipline/sources/__init__.py | 3 +- zipline/sources/data_frame_source.py | 69 ++++++++++++++++++++++++++++ zipline/utils/factory.py | 25 +++++++++- 6 files changed, 129 insertions(+), 7 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index f028c9a7..d757c04a 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -20,7 +20,9 @@ import numpy as np from zipline.utils.test_utils import setup_logger import zipline.utils.factory as factory from zipline.test_algorithms import TestRegisterTransformAlgorithm -from zipline.sources import SpecificEquityTrades, DataFrameSource +from zipline.sources import (SpecificEquityTrades, + DataFrameSource, + DataPanelSource) from zipline.transforms import MovingAverage @@ -42,6 +44,9 @@ class TestTransformAlgorithm(TestCase): self.df_source, self.df = \ factory.create_test_df_source(self.trading_environment) + self.panel_source, self.panel = \ + factory.create_test_panel_source(self.trading_environment) + def test_source_as_input(self): algo = TestRegisterTransformAlgorithm(sids=[133]) algo.run(self.source) @@ -64,6 +69,11 @@ class TestTransformAlgorithm(TestCase): algo.run(self.df) assert isinstance(algo.sources[0], DataFrameSource) + def test_panel_as_input(self): + algo = TestRegisterTransformAlgorithm(sids=[0, 1]) + algo.run(self.panel) + assert isinstance(algo.sources[0], DataPanelSource) + def test_run_twice(self): algo = TestRegisterTransformAlgorithm(sids=[0, 1]) res1 = algo.run(self.df) diff --git a/tests/test_sources.py b/tests/test_sources.py index 837ee459..e316a5d8 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -12,6 +12,7 @@ # 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. +import pandas as pd from unittest import TestCase @@ -20,8 +21,10 @@ from zipline.sources import DataFrameSource class TestDataFrameSource(TestCase): - def test_streaming_of_df(self): + def test_df_source(self): source, df = factory.create_test_df_source() + assert isinstance(source.start, pd.lib.Timestamp) + assert isinstance(source.end, pd.lib.Timestamp) for expected_dt, expected_price in df.iterrows(): sid0 = source.next() @@ -29,8 +32,23 @@ class TestDataFrameSource(TestCase): assert expected_dt == sid0.dt assert expected_price[0] == sid0.price - def test_sid_filtering(self): + def test_df_sid_filtering(self): _, df = factory.create_test_df_source() source = DataFrameSource(df, sids=[0]) assert 1 not in [event.sid for event in source], \ "DataFrameSource should only stream selected sid 0, not sid 1." + + def test_panel_source(self): + source, panel = factory.create_test_panel_source() + assert isinstance(source.start, pd.lib.Timestamp) + assert isinstance(source.end, pd.lib.Timestamp) + for event in source: + assert 'sid' in event + assert 'arbitrary' in event + assert 'volume' in event + assert 'price' in event + assert event['arbitrary'] == 1. + assert event['volume'] == 1000 + assert event['sid'] == 0 + assert isinstance(event['volume'], int) + assert isinstance(event['arbitrary'], float) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 736b0425..dda384fb 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -23,7 +23,7 @@ from datetime import datetime from itertools import groupby from operator import attrgetter -from zipline.sources import DataFrameSource +from zipline.sources import DataFrameSource, DataPanelSource from zipline.utils.factory import create_trading_environment from zipline.transforms.utils import StatefulTransform from zipline.finance.slippage import ( @@ -153,9 +153,10 @@ class TradingAlgorithm(object): """When providing a list of sources, \ start and end date have to be specified.""" elif isinstance(source, pd.DataFrame): - assert isinstance(source.index, pd.tseries.index.DatetimeIndex) # if DataFrame provided, wrap in DataFrameSource source = DataFrameSource(source) + elif isinstance(source, pd.Panel): + source = DataPanelSource(source) # If values not set, try to extract from source. if start is None: diff --git a/zipline/sources/__init__.py b/zipline/sources/__init__.py index e22a4c1e..2d853b4a 100644 --- a/zipline/sources/__init__.py +++ b/zipline/sources/__init__.py @@ -1,7 +1,8 @@ -from zipline.sources.data_frame_source import DataFrameSource +from zipline.sources.data_frame_source import DataFrameSource, DataPanelSource from zipline.sources.test_source import SpecificEquityTrades __all__ = [ 'DataFrameSource', + 'DataPanelSource', 'SpecificEquityTrades' ] diff --git a/zipline/sources/data_frame_source.py b/zipline/sources/data_frame_source.py index f9c43c07..57c0a9e1 100644 --- a/zipline/sources/data_frame_source.py +++ b/zipline/sources/data_frame_source.py @@ -81,3 +81,72 @@ class DataFrameSource(DataSource): if not self._raw_data: self._raw_data = self.raw_data_gen() return self._raw_data + + +class DataPanelSource(DataSource): + """ + Yields all events in event_list that match the given sid_filter. + If no event_list is specified, generates an internal stream of events + to filter. Returns all events if filter is None. + + Configuration options: + + sids : list of values representing simulated internal sids + start : start date + delta : timedelta between internal events + filter : filter to remove the sids + """ + + def __init__(self, data, **kwargs): + assert isinstance(data.major_axis, pd.tseries.index.DatetimeIndex) + + self.data = data + # Unpack config dictionary with default values. + self.sids = kwargs.get('sids', data.items) + self.start = kwargs.get('start', data.major_axis[0]) + self.end = kwargs.get('end', data.major_axis[-1]) + + # Hash_value for downstream sorting. + self.arg_string = hash_args(data, **kwargs) + + self._raw_data = None + + @property + def mapping(self): + mapping = { + 'dt': (lambda x: x, 'dt'), + 'sid': (lambda x: x, 'sid'), + 'price': (float, 'price'), + 'volume': (int, 'volume'), + } + + # Add additional fields. + for field_name in self.data.minor_axis: + if field_name in ['price', 'volume', 'dt', 'sid']: + continue + mapping[field_name] = (lambda x: x, field_name) + + return mapping + + @property + def instance_hash(self): + return self.arg_string + + def raw_data_gen(self): + for sid, dataframe in self.data.iteritems(): + for dt, series in dataframe.iterrows(): + if sid in self.sids: + event = { + 'dt': dt, + 'sid': sid, + } + for field_name, value in series.iteritems(): + event[field_name] = value + + yield event + + @property + def raw_data(self): + if not self._raw_data: + self._raw_data = self.raw_data_gen() + return self._raw_data diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index c203c914..c518d483 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -31,7 +31,9 @@ from datetime import datetime, timedelta import zipline.finance.risk as risk from zipline.utils.date_utils import tuple_to_date from zipline.utils.protocol_utils import ndict -from zipline.sources import SpecificEquityTrades, DataFrameSource +from zipline.sources import (SpecificEquityTrades, + DataFrameSource, + DataPanelSource) from zipline.gens.utils import create_trade from zipline.finance.trading import TradingEnvironment from zipline.data.loader import ( @@ -288,6 +290,27 @@ def create_test_df_source(trading_calendar=None): return DataFrameSource(df), df +def create_test_panel_source(trading_calendar=None): + start = trading_calendar.first_open \ + if trading_calendar else pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc) + + end = trading_calendar.last_close \ + if trading_calendar else pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc) + + index = pd.DatetimeIndex(start=start, end=end, freq=pd.datetools.day) + price = np.arange(0, len(index)) + volume = np.ones(len(index)) * 1000 + arbitrary = np.ones(len(index)) + + df = pd.DataFrame({'price': price, + 'volume': volume, + 'arbitrary': arbitrary}, + index=index) + panel = pd.Panel.from_dict({0: df}) + + return DataPanelSource(panel), panel + + def load_from_yahoo(indexes=None, stocks=None, start=None, end=None): """Load closing prices from yahoo finance.