Merge pull request #120 from quantopian/batch_window

Batch transform, new algorithm base class, new DataFrameSource

tests/test_algorithm.py

@@ -0,0 +1,57 @@
+from unittest2 import TestCase
+from datetime import timedelta
+
+from zipline.utils.test_utils import setup_logger
+import zipline.utils.factory as factory
+from zipline.test_algorithms import TestRegisterTransformAlgorithm
+from zipline.gens.tradegens import SpecificEquityTrades, DataFrameSource
+from zipline.gens.mavg import MovingAverage
+
+class TestTransformAlgorithm(TestCase):
+    def setUp(self):
+        setup_logger(self)
+        self.trading_environment = factory.create_trading_environment()
+        setup_logger(self)
+
+        trade_history = factory.create_trade_history(
+            133,
+            [10.0, 10.0, 11.0, 11.0],
+            [100, 100, 100, 300],
+            timedelta(days=1),
+            self.trading_environment
+        )
+        self.source = SpecificEquityTrades(event_list=trade_history)
+
+        self.df_source, self.df = factory.create_test_df_source()
+
+    def test_source_as_input(self):
+        algo = TestRegisterTransformAlgorithm(sids=[133])
+        algo.run(self.source)
+        self.assertEqual(len(algo.sources), 1)
+        assert isinstance(algo.sources[0], SpecificEquityTrades)
+
+    def test_multi_source_as_input_no_start_end(self):
+        algo = TestRegisterTransformAlgorithm(sids=[133])
+        with self.assertRaises(AssertionError):
+            algo.run([self.source, self.df_source])
+
+    def test_multi_source_as_input(self):
+        algo = TestRegisterTransformAlgorithm(sids=[0, 1, 133])
+        algo.run([self.source, self.df_source], start=self.df.index[0], end=self.df.index[-1])
+        self.assertEqual(len(algo.sources), 2)
+
+    def test_df_as_input(self):
+        algo = TestRegisterTransformAlgorithm(sids=[0, 1])
+        algo.run(self.df)
+        assert isinstance(algo.sources[0], DataFrameSource)
+
+    def test_transform_registered(self):
+        algo = TestRegisterTransformAlgorithm(sids=[133])
+        algo.run(self.source)
+        assert algo.get_sid_filter() == algo.sids == [133]
+        assert 'mavg' in algo.registered_transforms
+        assert algo.registered_transforms['mavg']['args'] == (['price'],)
+        assert algo.registered_transforms['mavg']['kwargs'] == {'days': 2, 'market_aware': True}
+        assert algo.registered_transforms['mavg']['class'] is MovingAverage
+
+

tests/test_optimize.py

@@ -5,12 +5,7 @@ from collections import defaultdict
 
 import numpy as np
 
-from zipline.core.devsimulator import AddressAllocator
-# TODO: refactor the factory to use generators
-# from zipline.optimize.factory import create_predictable_zipline
-
-DEFAULT_TIMEOUT = 15 # seconds
-EXTENDED_TIMEOUT = 90
+from zipline.optimize.factory import create_predictable_zipline
 
 from zipline.utils.test_utils import setup_logger, teardown_logger
 
@@ -24,7 +19,7 @@ class TestUpDown(TestCase):
 
     def setUp(self):
         self.zipline_test_config = {
-            'sid'         : 133,
+            'sid'         : [0],
             'trade_count' : 5,
             'amplitude'   : 30,
             'base_price'  : 50
@@ -36,7 +31,6 @@ class TestUpDown(TestCase):
         teardown_logger(self)
 
     @skip
-    @timed(DEFAULT_TIMEOUT)
     def test_source_and_orders(self):
         """verify that UpDownSource is having the correct
         behavior and that BuySellAlgorithm places the buy/sell
@@ -44,17 +38,17 @@ class TestUpDown(TestCase):
         UpDownSource and BuySellAlgorithm interact correctly."
 
         """
-        zipline, config = create_predictable_zipline(
+
+        algo, config = create_predictable_zipline(
             self.zipline_test_config,
-            offset=0,
-            simulate=False
+            offset=0
         )
 
         #extract arguments
         base_price = self.zipline_test_config['base_price']
         amplitude = self.zipline_test_config['amplitude']
 
-        prices = np.array([event.price for event in config['trade_source'].event_list])
+        prices = config['trade_source'][0].values
         max_price_idx = np.where(prices==prices.max())[0]
         min_price_idx = np.where(prices==prices.min())[0]
         self.assertTrue(np.all(max_price_idx % 2 == 1),
@@ -70,9 +64,9 @@ class TestUpDown(TestCase):
             "Minimum price does not equal expected maximum price."
         )
 
-        zipline.simulate(blocking=True)
+        stats = algo.run(config['trade_source'])
 
-        algo = config['algorithm']
+        self.assertTrue(len(stats) != 0)
 
         orders = np.asarray(algo.orders)
         max_order_idx = np.where(orders==orders.max())[0]
@@ -108,12 +102,15 @@ class TestUpDown(TestCase):
         compound_returns = np.empty(len(test_offsets))
         ziplines = []
         for i, offset in enumerate(test_offsets):
-            zipline, config = create_predictable_zipline(
+            algo, config = create_predictable_zipline(
                 self.zipline_test_config,
                 offset=offset,
             )
-            ziplines.append(zipline)
-            compound_returns[i] = zipline.get_cumulative_performance()['returns']
+            results = algo.run(config['trade_source'])
+            ziplines.append(algo)
+
+            compound_returns[i] = results.returns.sum()
+
 
         self.assertTrue(np.all(compound_returns[supposed_max] > compound_returns[np.logical_not(supposed_max)]),
             "Maximum compound returns are not where they are supposed to be."

tests/test_sources.py

@@ -0,0 +1,22 @@
+from unittest2 import TestCase
+
+import zipline.utils.factory as factory
+from zipline.gens.tradegens import DataFrameSource
+
+class TestDataFrameSource(TestCase):
+    def test_streaming_of_df(self):
+        source, df = factory.create_test_df_source()
+
+        for expected_dt, expected_price in df.iterrows():
+            sid0 = source.next()
+            sid1 = source.next()
+
+            assert expected_dt == sid0.dt == sid1.dt
+            assert expected_price[0] == sid0.price
+            assert expected_price[1] == sid1.price
+
+    def test_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."

tests/test_transforms.py

@@ -1,5 +1,5 @@
 import pytz
-import numpy
+import numpy as np
 
 from datetime import timedelta, datetime
 from unittest2 import TestCase
@@ -15,9 +15,10 @@ from zipline.gens.vwap import VWAP
 from zipline.gens.mavg import MovingAverage
 from zipline.gens.stddev import MovingStandardDev
 from zipline.gens.returns import Returns
-
 import zipline.utils.factory as factory
 
+from zipline.test_algorithms import BatchTransformAlgorithm
+
 def to_dt(msg):
     return ndict({'dt': msg})
 
@@ -42,26 +43,26 @@ class EventWindowTestCase(TestCase):
 
     def setUp(self):
         setup_logger(self)
-        
+
         self.monday = datetime(2012, 7, 9, 16, tzinfo=pytz.utc)
-        self.eleven_normal_days = [self.monday + i*timedelta(days=1) 
+        self.eleven_normal_days = [self.monday + i*timedelta(days=1)
                                    for i in xrange(11)]
 
         # Modify the end of the period slightly to exercise the
         # incomplete day logic.
         self.eleven_normal_days[-1] -= timedelta(minutes = 1)
         self.eleven_normal_days.append(self.monday+timedelta(days=11,seconds=1))
-        
+
         # Second set of dates to test holiday handling.
         self.jul4_monday = datetime(2012, 7, 2, 16, tzinfo=pytz.utc)
         self.week_of_jul4 = [self.jul4_monday + i*timedelta(days=1)
                                    for i in xrange(5)]
 
     def test_event_window_with_timedelta(self):
-        
+
         # Keep all events within a 5 minute window.
         window = NoopEventWindow(
-            market_aware = False, 
+            market_aware = False,
             delta = timedelta(minutes = 5),
             days = None
         )
@@ -91,7 +92,7 @@ class EventWindowTestCase(TestCase):
 
     def test_market_aware_window_normal_week(self):
         window = NoopEventWindow(
-            market_aware = True, 
+            market_aware = True,
             delta = None,
             days = 3
         )
@@ -102,7 +103,7 @@ class EventWindowTestCase(TestCase):
             window.update(event)
             # Record the length of the window after each event.
             lengths.append(len(window.ticks))
-        
+
         # The window stretches out during the weekend because we wait
         # to drop events until the weekend ends. The last window is
         # briefly longer because it doesn't complete a full day.  The
@@ -113,7 +114,7 @@ class EventWindowTestCase(TestCase):
 
     def test_market_aware_window_holiday(self):
         window = NoopEventWindow(
-            market_aware = True, 
+            market_aware = True,
             delta = None,
             days = 2
         )
@@ -125,11 +126,11 @@ class EventWindowTestCase(TestCase):
             window.update(event)
             # Record the length of the window after each event.
             lengths.append(len(window.ticks))
-        
+
         assert lengths == [1, 2, 3, 3, 2]
         assert window.added == events
         assert window.removed == events[:-2]
-        
+
     def tearDown(self):
         setup_logger(self)
 
@@ -186,7 +187,7 @@ class FinanceTransformsTestCase(TestCase):
         expected = [0.0, 0.0, 0.1, 0.0]
 
         assert tnfm_vals == expected
-        
+
         # Two-day returns.  An extra kink here is that the
         # factory will automatically skip a weekend for the
         # last event. Results shouldn't notice this blip.
@@ -222,12 +223,12 @@ class FinanceTransformsTestCase(TestCase):
             fields = ['price', 'volume'],
             delta = timedelta(days = 2),
         )
-        
+
         transformed = list(mavg.transform(self.source))
         # Output values.
         tnfm_prices = [message.tnfm_value.price for message in transformed]
         tnfm_volumes = [message.tnfm_value.volume for message in transformed]
-        
+
         # "Hand-calculated" values
         expected_prices = [
             ((10.0) / 1.0),
@@ -267,16 +268,16 @@ class FinanceTransformsTestCase(TestCase):
         transformed = list(stddev.transform(self.source))
 
         vals = [message.tnfm_value for message in transformed]
-        
+
         expected = [
             None,
-            numpy.std([10.0, 15.0], ddof = 1),
-            numpy.std([10.0, 15.0, 13.0], ddof = 1),
-            numpy.std([15.0, 13.0, 12.0], ddof = 1),
+            np.std([10.0, 15.0], ddof = 1),
+            np.std([10.0, 15.0, 13.0], ddof = 1),
+            np.std([15.0, 13.0, 12.0], ddof = 1),
             ]
 
-        # numpy has odd rounding behavior, cf.
-        # http://docs.scipy.org/doc/numpy/reference/generated/numpy.std.html
+        # np has odd rounding behavior, cf.
+        # http://docs.scipy.org/doc/np/reference/generated/np.std.html
         for v1, v2 in zip(vals, expected):
 
             if v1 == None:
@@ -285,8 +286,23 @@ class FinanceTransformsTestCase(TestCase):
             assert round(v1, 5) == round(v2, 5)
 
 
+############################################################
+# Test BatchTransform
 
-        
-        
-        
+class BatchTransformTestCase(TestCase):
+    def setUp(self):
+        setup_logger(self)
+        self.source, self.df = factory.create_test_df_source()
+
+    def test_batch_inherit(self):
+        algo = BatchTransformAlgorithm(sids=[0, 1])
+        algo.run(self.source)
+
+        assert algo.history_class[:2] == algo.history_decorator[:2] == [None, None], "First two iterations should return None"
+
+        # test overloaded class
+        for test_history in [algo.history_class, algo.history_decorator]:
+            self.assertTrue(np.all(test_history[2].values.flatten() == range(4, 10)))
+            self.assertTrue(np.all(test_history[3].values.flatten() == range(4, 10)))
+            self.assertTrue(np.all(test_history[4].values.flatten() == range(6, 14)))
 

zipline/algorithm.py

@@ -0,0 +1,184 @@
+import pandas as pd
+import numpy as np
+
+from zipline.gens.tradegens import DataFrameSource
+from zipline.utils.factory import create_trading_environment
+from zipline.gens.transform import StatefulTransform
+from zipline.lines import SimulatedTrading
+from zipline.finance.slippage import FixedSlippage
+
+
+class TradingAlgorithm(object):
+    """Base class for trading algorithms. Inherit and overload
+    initialize() and handle_data(data).
+
+    A new algorithm could look like this:
+    ```
+    class MyAlgo(TradingAlgorithm):
+        def initialize(amount):
+            self.amount = amount
+
+        def handle_data(data):
+            sid = self.sids[0]
+            self.order(sid, amount)
+    ```
+    To then to run this algorithm:
+
+    >>> my_algo = MyAlgo(100, sids=[0])
+    >>> stats = my_algo.run(data)
+
+    """
+    def __init__(self, sids, *args, **kwargs):
+        """
+        Initialize sids and other state variables.
+
+        Calls user-defined initialize and forwarding *args and **kwargs.
+        """
+        self.sids = sids
+        self.done = False
+        self.order = None
+        self.frame_count = 0
+        self.portfolio = None
+
+        self.registered_transforms = {}
+
+        # call to user-defined initialize method
+        self.initialize(*args, **kwargs)
+
+    def _create_simulator(self, start, end):
+        """
+        Create trading environment, transforms and SimulatedTrading object.
+
+        Gets called by self.run().
+        """
+        environment = create_trading_environment(start=start, end=end)
+
+        # Create transforms by wrapping them into StatefulTransforms
+        transforms = []
+        for namestring, trans_descr in self.registered_transforms.iteritems():
+            sf = StatefulTransform(
+                trans_descr['class'],
+                *trans_descr['args'],
+                **trans_descr['kwargs']
+            )
+            sf.namestring = namestring
+
+            transforms.append(sf)
+
+        # SimulatedTrading is the main class handling data streaming,
+        # application of transforms and calling of the user algo.
+        return SimulatedTrading(
+            self.sources,
+            transforms,
+            self,
+            environment,
+            FixedSlippage()
+        )
+
+    def run(self, source, start=None, end=None):
+        """Run the algorithm.
+
+        :Arguments:
+            source : can be either:
+                     - pandas.DataFrame
+                     - zipline source
+                     - list of zipline sources
+
+               If pandas.DataFrame is provided, it must have the
+               following structure:
+               * column names must consist of ints representing the
+                 different sids
+               * index must be DatetimeIndex
+               * array contents should be price info.
+
+        :Returns:
+            daily_stats : pandas.DataFrame
+              Daily performance metrics such as returns, alpha etc.
+
+        """
+        if isinstance(source, (list, tuple)):
+            assert start is not None and end is not None, \
+            "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, sids=self.sids)
+
+        # If values not set, try to extract from source.
+        if start is None:
+            start = source.start
+        if end is None:
+            end = source.end
+
+        if not isinstance(source, (list, tuple)):
+            self.sources = [source]
+        else:
+            self.sources = source
+
+        # create transforms and zipline
+        self.simulated_trading = self._create_simulator(start=start, end=end)
+
+        # loop through simulated_trading, each iteration returns a
+        # perf ndict
+        perfs = list(self.simulated_trading)
+
+        # convert perf ndict to pandas dataframe
+        daily_stats = self._create_daily_stats(perfs)
+
+        return daily_stats
+
+
+    def _create_daily_stats(self, perfs):
+        # create daily and cumulative stats dataframe
+        daily_perfs = []
+        cum_perfs = []
+        for perf in perfs:
+            if 'daily_perf' in perf:
+                daily_perfs.append(perf['daily_perf'])
+            else:
+                cum_perfs.append(perf)
+
+        daily_dts = [np.datetime64(perf['period_close'], utc=True) for perf in daily_perfs]
+        daily_stats = pd.DataFrame(daily_perfs, index=daily_dts)
+
+        return daily_stats
+
+
+    def add_transform(self, transform_class, tag, *args, **kwargs):
+        """Add a single-sid, sequential transform to the model.
+
+        :Arguments:
+            transform_class : class
+                Which transform to use. E.g. mavg.
+            tag : str
+                How to name the transform. Can later be access via:
+                data[sid].tag()
+
+        Extra args and kwargs will be forwarded to the transform
+        instantiation.
+
+        """
+        self.registered_transforms[tag] = {'class': transform_class,
+                                           'args': args,
+                                           'kwargs': kwargs}
+
+    def set_portfolio(self, portfolio):
+        self.portfolio = portfolio
+
+    def set_order(self, order_callable):
+        self.order = order_callable
+
+    def get_sid_filter(self):
+        return self.sids
+
+    def set_logger(self, logger):
+        self.logger = logger
+
+    def initialize(self, *args, **kwargs):
+        pass
+
+    def set_slippage_override(self, slippage_callable):
+        pass
+
+
+

zipline/gens/composites.py

@@ -1,10 +1,7 @@
-from itertools import tee, chain
+from itertools import chain
 
 from zipline.gens.utils import roundrobin, done_message
 from zipline.gens.sort import date_sort
-from zipline.gens.merge import merge
-from zipline.gens.transform import StatefulTransform
-
 
 def date_sorted_sources(*sources):
     """
@@ -14,7 +11,7 @@ def date_sorted_sources(*sources):
 
     for source in sources:
         assert iter(source), "Source %s not iterable" % source
-        assert 'get_hash' in source.__class__.__dict__, "No get_hash"
+        assert hasattr(source, 'get_hash'), "No get_hash"
 
     # Get name hashes to pass to date_sort.
     names = [source.get_hash() for source in sources]
@@ -29,46 +26,6 @@ def date_sorted_sources(*sources):
 
     return date_sort(stream_in, names)
 
-
-def merged_transforms(sorted_stream, *transforms):
-    """
-    A generator that takes the expected output of a date_sort, pipes
-    it through a given set of transforms, and runs the results
-    through a merge to output a unified stream. tnfms should be a
-    list of pointers to generator functions. tnfm_args should be a
-    list of tuples, representing the arguments to be passed to each
-    transform.  tnfm_kwargs should be a list of dictionaries
-    representing keyword arguments to each transform.
-    """
-    for transform in transforms:
-        assert isinstance(transform, StatefulTransform)
-        transform.merged = True
-        transform.sequential = False
-
-    # Generate expected hashes for each transform
-    namestrings = [tnfm.get_hash() for tnfm in transforms]
-
-    # Create a copy of the stream for each transform.
-    split = tee(sorted_stream, len(transforms))
-
-    # Package a stream copy with each StatefulTransform instance.
-    bundles = zip(transforms, split)
-
-    # Convert the copies into transform streams.
-    tnfm_gens = [tnfm.transform(stream) for tnfm, stream in bundles]
-
-    # Roundrobin the outputs of our transforms to create a single flat
-    # stream.
-    to_merge = roundrobin(tnfm_gens, namestrings)
-
-    # Pipe the stream into merge.
-    merged = merge(to_merge, namestrings)
-
-    dt_aliased = alias_dt(merged)
-    # Return the merged events.
-    return add_done(dt_aliased)
-
-
 def sequential_transforms(stream_in, *transforms):
     """
     Apply each transform in transforms sequentially to each event in stream_in.
@@ -87,6 +44,7 @@ def sequential_transforms(stream_in, *transforms):
                         transforms,
                         stream_in)
 
+
     dt_aliased = alias_dt(stream_out)
     return add_done(dt_aliased)
 

zipline/gens/tradegens.py

@@ -7,7 +7,12 @@ import pytz
 
 from itertools import chain, cycle, ifilter, izip, repeat
 from datetime import datetime, timedelta
+import pandas as pd
+from copy import copy
+import numpy as np
 
+from zipline.protocol import DATASOURCE_TYPE
+from zipline.utils import ndict
 from zipline.gens.utils import hash_args, create_trade
 
 def date_gen(start = datetime(2006, 6, 6, 12, tzinfo=pytz.utc),
@@ -73,17 +78,31 @@ class SpecificEquityTrades(object):
         # We shouldn't get any positional arguments.
         assert len(args) == 0
 
-        # Unpack config dictionary with default values.
-        self.count = kwargs.get('count', 500)
-        self.sids = kwargs.get('sids', [1, 2])
-        self.start = kwargs.get('start', datetime(2008, 6, 6, 15, tzinfo = pytz.utc))
-        self.delta = kwargs.get('delta', timedelta(minutes = 1))
-        self.concurrent = kwargs.get('concurrent', False)
-
         # Default to None for event_list and filter.
         self.event_list = kwargs.get('event_list')
         self.filter = kwargs.get('filter')
 
+        if self.event_list is not None:
+            # If event_list is provided, extract parameters from there
+            # This isn't really clean and ultimately I think this
+            # class should serve a single purpose (either take an
+            # event_list or autocreate events).
+            self.count = kwargs.get('count', len(self.event_list))
+            self.sids = kwargs.get('sids', np.unique([event.sid for event in self.event_list]).tolist())
+            self.start = kwargs.get('start', self.event_list[0].dt)
+            self.end = kwargs.get('start', self.event_list[-1].dt)
+            self.delta = kwargs.get('delta', self.event_list[1].dt - self.event_list[0].dt)
+            self.concurrent = kwargs.get('concurrent', False)
+
+        else:
+            # Unpack config dictionary with default values.
+            self.count = kwargs.get('count', 500)
+            self.sids = kwargs.get('sids', [1, 2])
+            self.start = kwargs.get('start', datetime(2008, 6, 6, 15, tzinfo = pytz.utc))
+            self.delta = kwargs.get('delta', timedelta(minutes = 1))
+            self.concurrent = kwargs.get('concurrent', False)
+
+
         # Hash_value for downstream sorting.
         self.arg_string = hash_args(*args, **kwargs)
 
@@ -137,6 +156,7 @@ class SpecificEquityTrades(object):
             volumes = mock_volumes(self.count)
 
             sids = cycle(self.sids)
+
             # Combine the iterators into a single iterator of arguments
             arg_gen = izip(sids, prices, volumes, dates)
 
@@ -157,33 +177,54 @@ class SpecificEquityTrades(object):
         return filtered
 
 
-# !!!!!!! Deprecated for now !!!!!!!!!
+class DataFrameSource(SpecificEquityTrades):
+    """
+    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.
 
-def RandomEquityTrades(object):
+    Configuration options:
 
-    def __init__(self):
-        # We shouldn't get any positional args.
-        assert args == ()
+    count  : integer representing number of trades
+    sids   : list of values representing simulated internal sids
+    start  : start date
+    delta  : timedelta between internal events
+    filter : filter to remove the sids
+    """
 
-        self.count = config.get('count', 500)
-        self.sids = config.get('sids', [1,2])
-        self.filter = config.get('filter')
+    def __init__(self, data, **kwargs):
+        assert isinstance(data.index, pd.tseries.index.DatetimeIndex)
 
-        dates = fuzzy_dates(count)
-        prices = mock_prices(count, rand = True)
-        volumes = mock_volumes(count, rand = True)
-        sids = cycle(sids)
+        self.data = data
+        # Unpack config dictionary with default values.
+        self.count = kwargs.get('count', len(data))
+        self.sids = kwargs.get('sids', data.columns)
+        self.start = kwargs.get('start', data.index[0])
+        self.end = kwargs.get('end', data.index[-1])
+        self.delta = kwargs.get('delta', data.index[1]-data.index[0])
 
-    arg_gen = izip(sids, prices, volumes, dates)
+        # Hash_value for downstream sorting.
+        self.arg_string = hash_args(data, **kwargs)
 
-    unfiltered = (create_trade(*args) for args in arg_gen)
+        self.generator = self.create_fresh_generator()
 
-    if filter:
-        filtered = ifilter(lambda event: event.sid in filter, unfiltered)
-    else:
-        filtered = unfiltered
-    return filtered
+    def create_fresh_generator(self):
+        def _generator(df=self.data):
+            for dt, series in df.iterrows():
+                if (dt < self.start) or (dt > self.end):
+                    continue
+                event = {'dt': dt,
+                         'source_id': self.get_hash(),
+                         'type': DATASOURCE_TYPE.TRADE
+                }
 
-# if __name__ == "__main__":
-#     import nose.tools; nose.tools.set_trace()
-#     trades = SpecificEquityTrades(filter = [1])
+                for sid, price in series.iterkv():
+                    event = copy(event)
+                    event['sid'] = sid
+                    event['price'] = price
+
+                    yield ndict(event)
+
+        # Return the filtered event stream.
+        drop_sids = lambda x: x.sid in self.sids
+        return ifilter(drop_sids, _generator())

zipline/gens/tradesimulation.py

@@ -2,6 +2,7 @@ from logbook import Logger, Processor
 
 from datetime import datetime
 from itertools import groupby
+from operator import attrgetter
 
 from zipline import ndict
 from zipline.utils.timeout import Heartbeat, Timeout
@@ -216,62 +217,61 @@ class AlgorithmSimulator(object):
         # Capture any output of this generator to stdout and pipe it
         # to a logbook interface.  Also inject the current algo
         # snapshot time to any log record generated.
-        with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''):
+        #with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''):
 
-            # Call user's initialize method with a timeout.
-            with Timeout(INIT_TIMEOUT, message="Call to initialize timed out"):
-                self.algo.initialize()
+        # Call user's initialize method with a timeout.
+        with Timeout(INIT_TIMEOUT, message="Call to initialize timed out"):
+            self.algo.initialize()
 
-            # Group together events with the same dt field. This depends on the
-            # events already being sorted.
-            for date, snapshot in groupby(stream_in, lambda e: e.dt):
+        # Group together events with the same dt field. This depends on the
+        # events already being sorted.
+        for date, snapshot in groupby(stream_in, attrgetter('dt')):
+            # Set the simulation date to be the first event we see.
+            # This should only occur once, at the start of the test.
+            if self.simulation_dt == None:
+                self.simulation_dt = date
 
-                # Set the simulation date to be the first event we see.
-                # This should only occur once, at the start of the test.
-                if self.simulation_dt == None:
-                    self.simulation_dt = date
+            # Done message has the risk report, so we yield before exiting.
+            if date == 'DONE':
+                for event in snapshot:
+                    yield event.perf_message
+                raise StopIteration()
 
-                # Done message has the risk report, so we yield before exiting.
-                if date == 'DONE':
-                    for event in snapshot:
+            # We're still in the warmup period.  Use the event to
+            # update our universe, but don't yield any perf messages,
+            # and don't send a snapshot to handle_data.
+            elif date < self.algo_start:
+                for event in snapshot:
+                    del event['perf_message']
+                    self.update_universe(event)
+
+            # The algo has taken so long to process events that
+            # its simulated time is later than the event time.
+            # Update the universe and yield any perf messages
+            # encountered, but don't call handle_data.
+            elif date < self.simulation_dt:
+                for event in snapshot:
+                    # Only yield if we have something interesting to say.
+                    if event.perf_message != None:
                         yield event.perf_message
-                    raise StopIteration()
+                    # Delete the message before updating so we don't send it
+                    # to the user.
+                    del event['perf_message']
+                    self.update_universe(event)
 
-                # We're still in the warmup period.  Use the event to
-                # update our universe, but don't yield any perf messages,
-                # and don't send a snapshot to handle_data.
-                elif date < self.algo_start:
-                    for event in snapshot:
-                        del event['perf_message']
-                        self.update_universe(event)
+            # Regular snapshot.  Update the universe and send a snapshot
+            # to handle data.
+            else:
+                for event in snapshot:
+                    # Only yield if we have something interesting to say.
+                    if event.perf_message != None:
+                        yield event.perf_message
+                    del event['perf_message']
 
-                # The algo has taken so long to process events that
-                # its simulated time is later than the event time.
-                # Update the universe and yield any perf messages
-                # encountered, but don't call handle_data.
-                elif date < self.simulation_dt:
-                    for event in snapshot:
-                        # Only yield if we have something interesting to say.
-                        if event.perf_message != None:
-                            yield event.perf_message
-                        # Delete the message before updating so we don't send it
-                        # to the user.
-                        del event['perf_message']
-                        self.update_universe(event)
+                    self.update_universe(event)
 
-                # Regular snapshot.  Update the universe and send a snapshot
-                # to handle data.
-                else:
-                    for event in snapshot:
-                        # Only yield if we have something interesting to say.
-                        if event.perf_message != None:
-                            yield event.perf_message
-                        del event['perf_message']
-
-                        self.update_universe(event)
-
-                    # Send the current state of the universe to the user's algo.
-                    self.simulate_snapshot(date)
+                # Send the current state of the universe to the user's algo.
+                self.simulate_snapshot(date)
 
     def update_universe(self, event):
         """

zipline/gens/transform.py

@@ -9,6 +9,8 @@ from datetime import datetime
 from collections import deque
 from abc import ABCMeta, abstractmethod
 
+import pandas as pd
+
 from zipline import ndict
 from zipline.utils.tradingcalendar import non_trading_days
 from zipline.gens.utils import assert_sort_unframe_protocol, hash_args
@@ -36,7 +38,7 @@ class TransformMeta(type):
     still recover an instance of a "raw" Foo by introspecting the
     resulting StatefulTransform's 'state' field.
     """
-    
+
     def __call__(cls, *args, **kwargs):
         return StatefulTransform(cls, *args, **kwargs)
 
@@ -53,19 +55,19 @@ class StatefulTransform(object):
     def __init__(self, tnfm_class, *args, **kwargs):
         assert isinstance(tnfm_class, (types.ObjectType, types.ClassType)), \
         "Stateful transform requires a class."
-        assert tnfm_class.__dict__.has_key('update'), \
+        assert hasattr(tnfm_class, 'update'), \
         "Stateful transform requires the class to have an update method"
 
         # Flag set inside the Passthrough transform class to signify special
         # behavior if we are being fed to merged_transforms.
-        self.passthrough = tnfm_class.__dict__.get('PASSTHROUGH', False)
-        
+        self.passthrough = hasattr(tnfm_class, 'PASSTHROUGH')
+
         # Flags specifying how to append the calculated value.
         # Merged is the default for ease of testing, but we use sequential
         # in production.
         self.sequential = False
         self.merged = True
-        
+
         # Create an instance of our transform class.
         if isinstance(tnfm_class, TransformMeta):
             # Classes derived TransformMeta have their __call__
@@ -104,12 +106,12 @@ class StatefulTransform(object):
                 continue
 
             assert_sort_unframe_protocol(message)
-            
+
             # This flag is set by by merged_transforms to ensure
             # isolation of messages.
             if self.merged:
                 message = deepcopy(message)
-            
+
             tnfm_value = self.state.update(message)
 
             # PASSTHROUGH flag means we want to keep all original
@@ -133,7 +135,7 @@ class StatefulTransform(object):
                 out_message.tnfm_value = tnfm_value
                 out_message.dt = message.dt
                 yield out_message
-            
+
             # Sequential flag should be used to add a single new
             # key-value pair to the event. The new key is this
             # transform's namestring, and its value is the value
@@ -147,9 +149,11 @@ class StatefulTransform(object):
                 out_message = message
                 out_message[self.namestring] = tnfm_value
                 yield out_message
-                
+
         log.info('Finished StatefulTransform [%s]' % self.get_hash())
 
+
+
 class EventWindow(object):
     """
     Abstract base class for transform classes that calculate iterative
@@ -171,13 +175,13 @@ class EventWindow(object):
     # Mark this as an abstract base class.
     __metaclass__ = ABCMeta
 
-    def __init__(self, market_aware, days = None, delta = None):
+    def __init__(self, market_aware, days=None, delta=None):
 
         self.market_aware = market_aware
         self.days = days
         self.delta = delta
 
-        self.ticks  = deque()
+        self.ticks = deque()
 
         # Market-aware mode only works with full-day windows.
         if self.market_aware:
@@ -213,12 +217,12 @@ class EventWindow(object):
         self.assert_well_formed(event)
 
         # Add new event and increment totals.
-        self.ticks.append(event)
+        self.ticks.append(deepcopy(event))
 
         # Subclasses should override handle_add to define behavior for
         # adding new ticks.
         self.handle_add(event)
-        
+
         if self.market_aware:
             self.add_new_holidays(event.dt)
 
@@ -229,14 +233,14 @@ class EventWindow(object):
         #                                |                    |
         #                                V                    V
         while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt):
-            
+
             # popleft removes and returns the oldest tick in self.ticks
             popped = self.ticks.popleft()
 
             # Subclasses should override handle_remove to define
             # behavior for removing ticks.
             self.handle_remove(popped)
-            
+
     def add_new_holidays(self, newest):
         # Add to our tracked window any untracked holidays that are
         # older than our newest event. (newest should always be
@@ -256,12 +260,13 @@ class EventWindow(object):
         calendar_dates_between = (newest.date() - oldest.date()).days
         holidays_between = len(self.cur_holidays)
         trading_days_between = calendar_dates_between - holidays_between
-        
+
         # "Put back" a day if oldest is earlier in its day than newest,
         # reflecting the fact that we haven't yet completed the last
         # day in the window.
         if oldest.time() > newest.time():
             trading_days_between -= 1
+
         return trading_days_between >= self.days
 
     def out_of_delta(self, oldest, newest):
@@ -277,3 +282,144 @@ class EventWindow(object):
             # Something is wrong if new event is older than previous.
             assert event.dt >= self.ticks[-1].dt, \
                 "Events arrived out of order in EventWindow: %s -> %s" % (event, self.ticks[0])
+
+
+class BatchTransform(EventWindow):
+    """Base class for batch transforms with a trailing window of
+    variable length. As opposed to pure EventWindows that get a stream
+    of events and are bound to a single SID, this class creates stream
+    of pandas DataFrames with each colum representing a sid.
+
+    There are two ways to create a new batch window:
+    (i) Inherit from BatchTransform and overload get_value(data).
+        E.g.:
+        ```
+        class MyBatchTransform(BatchTransform):
+            def get_value(self, data):
+               # compute difference between the means of sid 0 and sid 1
+               return data[0].mean() - data[1].mean()
+        ```
+
+    (ii) Use the batch_transform decorator.
+        E.g.:
+        ```
+        @batch_transform
+        def my_batch_transform(data):
+            return data[0].mean() - data[1].mean()
+
+        ```
+
+    In you algorithm you would then have to instantiate this in the initialize() method:
+    ```
+    self.my_batch_transform = MyBatchTransform()
+    ```
+
+    To then use it, inside of the algorithm handle_data(), call the
+    handle_data() of the BatchTransform and pass it the current event:
+    ```
+    result = self.my_batch_transform(data)
+    ```
+
+    """
+
+    def __init__(self, func=None, refresh_period=None, market_aware=True, delta=None, days=None, sids=None):
+        super(BatchTransform, self).__init__(market_aware, days=days, delta=delta)
+        if func is not None:
+            self.compute_transform_value = func
+        else:
+            self.compute_transform_value = self.get_value
+
+        self.sids = sids
+        self.refresh_period = refresh_period
+        self.days = days
+
+        self.full = False
+        self.last_refresh = None
+
+        self.updated = False
+        self.data = None
+
+    def handle_data(self, data):
+        """
+        New method to handle a data frame as sent to the algorithm's handle_data
+        method.
+        """
+        # extract dates
+        dts = [data[sid].datetime for sid in self.sids]
+        # we have to provide the event with a dt. This is only for
+        # checking if the event is outside the window or not so a
+        # couple of seconds shouldn't matter
+        data.dt = max(dts)
+
+        # append data frame to window. update() will call handle_add() and
+        # handle_remove() appropriately
+        self.update(data)
+
+        # return newly computed or cached value
+        return self.get_transform_value()
+
+    def handle_add(self, event):
+        if not self.last_refresh:
+            self.last_refresh = event.dt
+            return
+
+        age = event.dt - self.last_refresh
+        if age.days >= self.refresh_period:
+            # Create a pandas.Panel (i.e. 3d DataFrame) from the
+            # events in the current window.
+            #
+            # The resulting panel looks like this:
+            # index : field_name (e.g. price)
+            # major axis/rows : dt
+            # minor axis/colums : sid
+            #
+            # This Panel data structure ultimately gets passed to the
+            # user-overloaded get_value() method.
+            fields = {}
+            for field_name in ['price', 'volume']:
+                # Skip non-existant fields
+                if field_name not in self.ticks[0][self.sids[0]]:
+                    continue
+
+                values_per_sid = {}
+                for sid in self.sids:
+                    values_per_sid[sid] = pd.Series({tick[sid].dt: tick[sid][field_name] for tick in self.ticks})
+
+                # concatenate different sids into one df
+                fields[field_name] = pd.DataFrame.from_dict(values_per_sid)
+
+            self.data = pd.Panel.from_dict(fields, orient='items')
+
+            self.updated = True
+            self.last_refresh = event.dt
+        else:
+            self.updated = False
+
+    def handle_remove(self, event):
+        # since an event is expiring, we know the window is full
+        self.full = True
+
+    def get_value(self, *args, **kwargs):
+        raise NotImplementedError("Either overwrite get_value or provide a func argument.")
+
+    def get_transform_value(self, *args, **kwargs):
+        if self.data is None:
+            return None
+
+        if self.updated:
+            self.cached = self.compute_transform_value(self.data, *args, **kwargs)
+
+        return self.cached
+
+
+def batch_transform(func):
+    """Decorator function to use instead of inheriting from BatchTransform.
+    For an example on how to use this, see the doc string of BatchTransform.
+    """
+
+    def create_window(*args, **kwargs):
+        # passes the user defined function to BatchTransform which it
+        # will call instead of self.get_value()
+        return BatchTransform(*args, func=func, **kwargs)
+
+    return create_window

zipline/optimize/algorithms.py

@@ -1,4 +1,9 @@
-class BuySellAlgorithm(object):
+from logbook import Logger
+from zipline.algorithm import TradingAlgorithm
+
+logger = Logger('Algo')
+
+class BuySellAlgorithm(TradingAlgorithm):
     """Algorithm that buys and sells alternatingly. The amount for
     each order can be specified. In addition, an offset that will
     quadratically reduce the amount that will be bought can be
@@ -10,39 +15,19 @@ class BuySellAlgorithm(object):
 
     """
 
-    def __init__(self, sid, amount, offset):
-        self.sid = sid
+    def initialize(self, amount=100, offset=0):
         self.amount = amount
-        self.incr = 0
-        self.done = False
-        self.order = None
-        self.frame_count = 0
-        self.portfolio = None
         self.buy_or_sell = -1
         self.offset = offset
         self.orders = []
-        self.prices = []
 
-    def initialize(self):
-        pass
-
-    def set_order(self, order_callable):
-        self.order = order_callable
-
-    def set_portfolio(self, portfolio):
-        self.portfolio = portfolio
-
-    def handle_data(self, frame):
+    def handle_data(self, data):
         order_size = self.buy_or_sell * (self.amount - (self.offset**2))
-        self.order(self.sid, order_size)
+        self.order(self.sids[0], order_size)
+        logger.debug("ordering" + str(order_size))
 
         #sell next time around.
         self.buy_or_sell *= -1
 
         self.orders.append(order_size)
 
-        self.frame_count += 1
-        self.incr += 1
-
-    def get_sid_filter(self):
-        return [self.sid]

zipline/optimize/example.py

@@ -0,0 +1,159 @@
+# WARNING: This file is still work in progress and contains rather
+# random code snippets.
+
+import pandas as pd
+
+import numpy as np
+
+import matplotlib.pyplot as plt
+import cProfile
+from zipline.gens.mavg import MovingAverage
+from zipline.gens.cov import  CovTransform, cov
+from zipline.algorithm import TradingAlgorithm
+from zipline.gens.transform import BatchTransform, batch_transform
+
+@batch_transform
+def cov(data):
+    return data.price.cov()
+
+class DMA(TradingAlgorithm):
+    """Dual Moving Average algorithm.
+    """
+    def initialize(self, amount=100, short_window=20, long_window=40):
+        self.amount = amount
+        self.events = 0
+
+        self.invested = {}
+        for sid in self.sids:
+            self.invested[sid] = False
+
+        self.add_transform(MovingAverage, 'short_mavg', ['price'],
+                           market_aware=True,
+                           days=short_window)
+
+        self.add_transform(MovingAverage, 'long_mavg', ['price'],
+                           market_aware=True,
+                           days=long_window)
+
+        self.cov = cov(sids=self.sids, refresh_period=1, days=5)
+
+    def handle_data(self, data):
+        self.events += 1
+
+        cov = self.cov.handle_data(data)
+        print cov
+
+        for sid in self.sids:
+            # access transforms via their user-defined tag
+            if (data[sid].short_mavg['price'] > data[sid].long_mavg['price']) and not self.invested[sid]:
+                self.order(sid, self.amount)
+                self.invested[sid] = True
+            elif (data[sid].short_mavg['price'] < data[sid].long_mavg['price']) and self.invested[sid]:
+                self.order(sid, -self.amount)
+                self.invested[sid] = False
+
+
+def load_close_px(indexes=None, stocks=None):
+    from pandas.io.data import DataReader
+    import pytz
+
+    if indexes is None:
+        indexes = {'SPX' : '^GSPC'}
+    if stocks is None:
+        stocks = ['AAPL'] #, 'GE', 'IBM', 'MSFT', 'XOM', 'AA', 'JNJ', 'PEP']
+
+    #start = pd.datetime(1990, 1, 1)
+    start = pd.datetime(1990, 1, 1, 0, 0, 0, 0, pytz.utc)
+    end = pd.datetime(1992, 1, 1, 0, 0, 0, 0, pytz.utc) #pd.datetime.today()
+
+    data = {}
+
+    for stock in stocks:
+        print stock
+        stkd = DataReader(stock, 'yahoo', start, end).sort_index()
+        data[stock] = stkd
+
+    for name, ticker in indexes.iteritems():
+        print name
+        stkd = DataReader(ticker, 'yahoo', start, end).sort_index()
+        data[name] = stkd
+
+    #df = pd.DataFrame({key: d['Close'] for key, d in data.iteritems()})
+    df = pd.DataFrame({i: d['Close'] for i, d in enumerate(data.itervalues())})
+    df.index = df.index.tz_localize(pytz.utc)
+
+    return df
+
+
+def run((short_window, long_window)):
+    #data = pd.DataFrame.from_csv('SP500.csv')
+    #data = pd.DataFrame.from_csv('aapl.csv') #load_close_px()
+    data = load_close_px()
+    myalgo = DMA([0, 1], amount=100, short_window=short_window, long_window=long_window)
+    stats = myalgo.run(data)
+    stats['sw'] = short_window
+    stats['lw'] = long_window
+    return stats
+
+def explore_params():
+    sws, lws = np.mgrid[10:20:5, 10:20:5]
+
+    stats_all = map(run, zip(sws.flatten(), lws.flatten()))
+    stats = pd.concat(stats_all)
+    returns = stats.groupby(['sw', 'lw']).sum()
+
+    plt.contourf(sws, lws, returns.returns.reshape(sws.shape))
+    plt.xlabel('Short window length')
+    plt.ylabel('Long window length')
+    plt.savefig('DMA_contour.png')
+    plt.show()
+
+#stats = run((10, 50))
+
+def get_opt_holdings_qp(univ_rets, track_rets):
+    from cvxopt import matrix
+    from cvxopt.solvers import qp
+    # set up the QP for CVXOPT
+    # .5 x' P x + q'x
+    # P = 2 * R'R
+    # q = - 2 * bmk'R
+    R = univ_rets.values
+    b = track_rets.values
+    P = matrix(2 * np.dot(R.T, R))
+    q = matrix(-2 * np.dot(R.T, b))
+    result = qp(P, q)
+    if result['status'] != 'optimal':
+        raise Exception('optimum not reached by QP')
+    return pd.Series(np.array(result['x']).ravel(), index=univ_rets.columns)
+
+def opt_portfolio(cov, budget, min_return):
+    from cvxopt import matrix
+    from cvxopt.solvers import qp
+    n = len(cov)
+    cov = matrix(2 * cov)
+    q = matrix(np.zeros(n))
+
+    h = matrix(budget) # G*x < h
+    # coneqp
+    result = qp(cov, q, h=h)
+    if result['status'] != 'optimal':
+        raise Exception('optimum not reached by QP')
+
+    return pd.Series(np.array(result['x']).ravel())
+
+def calc_te(weights, univ_rets, track_rets):
+    port_rets = (univ_rets * weights).sum(1)
+    return (port_rets - track_rets).std()
+
+def plot_returns(port_returns, bmk_returns):
+    plt.figure()
+    cum_port = ((1 + port_returns).cumprod() - 1)
+    cum_bmk = ((1 + bmk_returns).cumprod() - 1)
+    # cum_port = port_returns.cumsum()
+    # cum_bmk = bmk_returns.cumsum()
+    cum_port.plot(label='Portfolio returns')
+    cum_bmk.plot(label='Benchmark')
+    plt.title('Portfolio performance')
+    plt.legend(loc='best')
+
+print run((10, 20))

zipline/optimize/factory.py

@@ -4,14 +4,12 @@ Factory functions to prepare useful data for optimize tests.
 Author: Thomas V. Wiecki (thomas.wiecki@gmail.com), 2012
 """
 from datetime import timedelta
-import pandas as pd
 
 import zipline.protocol as zp
 
 from zipline.utils.factory import get_next_trading_dt, create_trading_environment
-from zipline.finance.sources import SpecificEquityTrades
+from zipline.gens.tradegens import SpecificEquityTrades
 from zipline.optimize.algorithms import BuySellAlgorithm
-from zipline.lines import SimulatedTrading
 from zipline.finance.slippage import FixedSlippage
 
 from copy import copy
@@ -122,7 +120,7 @@ def create_predictable_zipline(config, offset=0, simulate=True):
                                         amplitude)
 
     if 'algorithm' not in config:
-        config['algorithm'] = BuySellAlgorithm(sid, 100, offset)
+        algorithm = BuySellAlgorithm(sids=[sid], amount=100, offset=offset)
 
     config['order_count'] = trade_count - 1
     config['trade_count'] = trade_count
@@ -131,9 +129,4 @@ def create_predictable_zipline(config, offset=0, simulate=True):
     config['slippage'] = FixedSlippage()
     config['devel'] = True
 
-    zipline = SimulatedTrading.create_test_zipline(**config)
-
-    if simulate:
-        zipline.simulate(blocking=True)
-
-    return zipline, config
+    return algorithm, config

zipline/test_algorithms.py

@@ -69,6 +69,7 @@ class TestAlgorithm():
         self.order = None
         self.frame_count = 0
         self.portfolio = None
+
         if sid_filter:
             self.sid_filter = sid_filter
         else:
@@ -99,7 +100,7 @@ class TestAlgorithm():
     def set_slippage_override(self, slippage_callable):
         pass
 
-    #
+
 class HeavyBuyAlgorithm():
     """
     This algorithm will send a specified number of orders, to allow unit tests
@@ -382,3 +383,49 @@ class TestLoggingAlgorithm():
 
     def set_slippage_override(self, slippage_callable):
         pass
+
+
+from datetime import timedelta
+from zipline.algorithm import TradingAlgorithm
+from zipline.gens.transform import BatchTransform, batch_transform
+from zipline.gens.mavg import MovingAverage
+
+class TestRegisterTransformAlgorithm(TradingAlgorithm):
+    def initialize(self):
+        self.add_transform(MovingAverage, 'mavg', ['price'],
+                           market_aware=True,
+                           days=2)
+
+    def handle_data(self, data):
+        pass
+
+class NoopBatchTransform(BatchTransform):
+    def get_value(self, data):
+        return data.price
+
+@batch_transform
+def noop_batch_decorator(data):
+    return data.price
+
+class BatchTransformAlgorithm(TradingAlgorithm):
+    def initialize(self, *args, **kwargs):
+        self.history_class = []
+        self.history_decorator = []
+        self.days = 3
+        self.noop_class = NoopBatchTransform(sids=[0, 1],
+                                          market_aware=False,
+                                          refresh_period=2,
+                                          delta=timedelta(days=self.days))
+
+        self.noop_decorator = noop_batch_decorator(sids=[0, 1],
+                                                   market_aware=False,
+                                                   refresh_period=2,
+                                                   delta=timedelta(days=self.days))
+
+    def handle_data(self, data):
+        window_class = self.noop_class.handle_data(data)
+        window_decorator = self.noop_decorator.handle_data(data)
+        self.history_class.append(window_class)
+        self.history_decorator.append(window_decorator)
+
+

zipline/utils/factory.py

@@ -8,14 +8,14 @@ import random
 from os.path import join, abspath, dirname
 from operator import attrgetter
 
+import pandas as pd
+import numpy as np
 from datetime import datetime, timedelta
-from zipline.utils.date_utils import tuple_to_date
-from zipline.utils.protocol_utils import ndict
 
 import zipline.finance.risk as risk
-
-from zipline.gens.tradegens import RandomEquityTrades
-from zipline.gens.tradegens import SpecificEquityTrades
+from zipline.utils.date_utils import tuple_to_date
+from zipline.utils.protocol_utils import ndict
+from zipline.gens.tradegens import SpecificEquityTrades, DataFrameSource
 from zipline.gens.utils import create_trade
 from zipline.finance.trading import TradingEnvironment
 
@@ -57,12 +57,15 @@ def load_market_data():
 
     return bm_returns, tr_curves
 
-def create_trading_environment(year=2006):
+def create_trading_environment(year=2006, start=None, end=None):
     """Construct a complete environment with reasonable defaults"""
     benchmark_returns, treasury_curves = load_market_data()
 
-    start = datetime(year, 1, 1, tzinfo=pytz.utc)
-    end   = datetime(year, 12, 31, tzinfo=pytz.utc)
+    if start is None:
+        start = datetime(year, 1, 1, tzinfo=pytz.utc)
+    if end is None:
+        end   = datetime(year, 12, 31, tzinfo=pytz.utc)
+
     trading_environment = TradingEnvironment(
         benchmark_returns,
         treasury_curves,
@@ -88,7 +91,6 @@ def create_trade_history(sid, prices, amounts, interval, trading_calendar):
     current = trading_calendar.first_open
 
     for price, amount in zip(prices, amounts):
-
         trade = create_trade(sid, price, amount, current)
         trades.append(trade)
         current = get_next_trading_dt(current, interval, trading_calendar)
@@ -233,3 +235,12 @@ def create_trade_source(sids, trade_count, trade_time_increment, trading_environ
     #trading_environment.period_end = trade_history[-1].dt
 
     return source
+
+def create_test_df_source():
+    start = pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc)
+    end = pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)
+    index = pd.DatetimeIndex(start=start, end=end, freq=pd.datetools.day)
+    x = np.arange(2., 14.).reshape((6, 2))
+    df = pd.DataFrame(x, index=index, columns=[0, 1])
+
+    return DataFrameSource(df), df