Merge pull request #846 from grundgruen/data_test

TST: tests removing of expired data and removes ffill in DataPanelSource

tests/test_algorithm.py

@@ -60,6 +60,7 @@ from zipline.test_algorithms import (
     TestTargetAlgorithm,
     TestTargetPercentAlgorithm,
     TestTargetValueAlgorithm,
+    TestRemoveDataAlgo,
     SetLongOnlyAlgorithm,
     SetAssetDateBoundsAlgorithm,
     SetMaxPositionSizeAlgorithm,
@@ -1977,3 +1978,43 @@ class TestTradingAlgorithm(TestCase):
                                 analyze=analyze)
         results = algo.run(self.panel)
         self.assertIs(results, self.perf_ref)
+
+
+class TestRemoveData(TestCase):
+    """
+    tests if futures data is removed after expiry
+    """
+    def setUp(self):
+        dt = pd.Timestamp('2015-01-02', tz='UTC')
+        env = TradingEnvironment()
+        ix = env.trading_days.get_loc(dt)
+
+        metadata = {0: {'symbol': 'X',
+                        'expiration_date': env.trading_days[ix + 5],
+                        'end_date': env.trading_days[ix + 6]},
+                    1: {'symbol': 'Y',
+                        'expiration_date': env.trading_days[ix + 7],
+                        'end_date': env.trading_days[ix + 8]}}
+
+        env.write_data(futures_data=metadata)
+
+        index_x = env.trading_days[ix:ix + 5]
+        data_x = pd.DataFrame([[1, 100], [2, 100], [3, 100], [4, 100],
+                               [5, 100]],
+                              index=index_x, columns=['price', 'volume'])
+        index_y = env.trading_days[ix:ix + 5].shift(2)
+        data_y = pd.DataFrame([[6, 100], [7, 100], [8, 100], [9, 100],
+                               [10, 100]],
+                              index=index_y, columns=['price', 'volume'])
+
+        pan = pd.Panel({0: data_x, 1: data_y})
+        self.source = DataPanelSource(pan)
+        self.algo = TestRemoveDataAlgo(env=env)
+
+    def test_remove_data(self):
+        self.algo.run(self.source)
+
+        expected_lengths = [1, 1, 2, 2, 2, 2, 1]
+        # initially only data for X should be sent and on the last day only
+        # data for Y should be sent since X is expired
+        np.testing.assert_array_equal(self.algo.data, expected_lengths)

tests/test_events_through_risk.py

@@ -15,6 +15,7 @@
 
 import unittest
 import datetime
+import pandas as pd
 import pytz
 
 import numpy as np
@@ -77,9 +78,9 @@ class TestEventsThroughRisk(unittest.TestCase):
 
         algo = BuyAndHoldAlgorithm(sim_params=sim_params, env=self.env)
 
-        first_date = datetime.datetime(2006, 1, 3, tzinfo=pytz.utc)
-        second_date = datetime.datetime(2006, 1, 4, tzinfo=pytz.utc)
-        third_date = datetime.datetime(2006, 1, 5, tzinfo=pytz.utc)
+        first_date = pd.Timestamp('2006-01-03', tz='UTC')
+        second_date = pd.Timestamp('2006-01-04', tz='UTC')
+        third_date = pd.Timestamp('2006-01-05', tz='UTC')
 
         trade_bar_data = [
             Event({

tests/test_sources.py

@@ -123,9 +123,7 @@ class TestDataFrameSource(TestCase):
         self.assertEqual(5, event.sid)
         event = next(source)
         self.assertEqual(4, event.sid)
-        event = next(source)
-        self.assertEqual(5, event.sid)
-        self.assertFalse(np.isnan(event.price))
+        self.assertRaises(StopIteration, next, source)
 
 
 class TestRandomWalkSource(TestCase):

zipline/assets/assets.py

@@ -446,9 +446,9 @@ class AssetFinder(object):
                 self.equities.c.share_class_symbol ==
                 share_class_symbol,
                 self.equities.c.start_date <= ad_value),
-            ).order_by(
-                self.equities.c.end_date.desc(),
-            ).execute().fetchall()
+        ).order_by(
+            self.equities.c.end_date.desc(),
+        ).execute().fetchall()
         return candidates
 
     def _get_best_candidate(self, candidates):
@@ -656,6 +656,26 @@ class AssetFinder(object):
         contracts = self.retrieve_futures_contracts(sids)
         return [contracts[sid] for sid in sids]
 
+    def lookup_expired_futures(self, start, end):
+        start = start.value
+        end = end.value
+
+        fc_cols = self.futures_contracts.c
+
+        nd = sa.func.nullif(fc_cols.notice_date, pd.tslib.iNaT)
+        ed = sa.func.nullif(fc_cols.expiration_date, pd.tslib.iNaT)
+        date = sa.func.coalesce(sa.func.min(nd, ed), ed, nd)
+
+        sids = list(map(
+            itemgetter('sid'),
+            sa.select((fc_cols.sid,)).where(
+                (date >= start) & (date < end)).order_by(
+                sa.func.coalesce(ed, nd).asc()
+            ).execute().fetchall()
+        ))
+
+        return sids
+
     @property
     def sids(self):
         return tuple(map(
@@ -904,6 +924,7 @@ class AssetFinderCachedEquities(AssetFinder):
     into memory and overrides the methods that lookup_symbol uses to look up
     those equities.
     """
+
     def __init__(self, engine):
         super(AssetFinderCachedEquities, self).__init__(engine)
         self.fuzzy_symbol_hashed_equities = {}

zipline/gens/tradesimulation.py

@@ -64,6 +64,7 @@ class AlgorithmSimulator(object):
         # We don't have a datetime for the current snapshot until we
         # receive a message.
         self.simulation_dt = None
+        self.previous_dt = self.algo_start
 
         # =============
         # Logging Setup
@@ -96,10 +97,19 @@ class AlgorithmSimulator(object):
             self._call_before_trading_start(mkt_open)
 
             for date, snapshot in stream_in:
-
+                expired_sids = self.env.asset_finder.lookup_expired_futures(
+                    start=self.previous_dt, end=date)
+                self.previous_dt = date
                 self.simulation_dt = date
                 self.on_dt_changed(date)
 
+                # removing expired futures
+                for sid in expired_sids:
+                    try:
+                        del self.current_data[sid]
+                    except KeyError:
+                        continue
+
                 # If 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.

zipline/sources/data_frame_source.py

@@ -114,7 +114,6 @@ class DataPanelSource(DataSource):
         # TODO is ffilling correct/necessary?
         # forward fill with volumes of 0
         self.data = data.fillna(value={'volume': 0})
-        self.data = self.data.fillna(method='ffill')
         # Unpack config dictionary with default values.
         self.start = kwargs.get('start', self.data.major_axis[0])
         self.end = kwargs.get('end', self.data.major_axis[-1])
@@ -153,8 +152,7 @@ class DataPanelSource(DataSource):
             df = self.data.major_xs(dt)
             for sid, series in df.iteritems():
                 # Skip SIDs that can not be forward filled
-                if np.isnan(series['price']) and \
-                   sid not in self.started_sids:
+                if np.isnan(series['price']):
                     continue
                 self.started_sids.add(sid)
 

zipline/test_algorithms.py

@@ -937,6 +937,16 @@ class InvalidOrderAlgorithm(TradingAlgorithm):
                                      style=style)
 
 
+class TestRemoveDataAlgo(TradingAlgorithm):
+    def initialize(self, *args, **kwargs):
+        self.data = np.zeros(7)
+        self.i = 0
+
+    def handle_data(self, data):
+        self.data[self.i] = len(data)
+        self.i += 1
+
+
 ##############################
 # Quantopian style algorithms