diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 2b765f28..d1f3ac76 100644 --- a/tests/test_algorithm.py +++ b/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) diff --git a/tests/test_events_through_risk.py b/tests/test_events_through_risk.py index b40ea1ca..91649931 100644 --- a/tests/test_events_through_risk.py +++ b/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({ diff --git a/tests/test_sources.py b/tests/test_sources.py index ff82e2fe..7c5c099e 100644 --- a/tests/test_sources.py +++ b/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): diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 13b457b4..149cb8c9 100644 --- a/zipline/assets/assets.py +++ b/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 = {} diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index e63ca853..f720333b 100644 --- a/zipline/gens/tradesimulation.py +++ b/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. diff --git a/zipline/sources/data_frame_source.py b/zipline/sources/data_frame_source.py index 50c0b9b7..eab04a61 100644 --- a/zipline/sources/data_frame_source.py +++ b/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) diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index a56d001f..ffe39226 100644 --- a/zipline/test_algorithms.py +++ b/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