Merge pull request #43 from quantopian/extend_df_source

ENH: Added DataPanelSource.

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)

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)

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:

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'
 ]

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

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.