From 5be63f36d527947033b42947a19ff72b14963468 Mon Sep 17 00:00:00 2001 From: dmichalowicz Date: Thu, 28 Jan 2016 10:16:33 -0500 Subject: [PATCH] ENH: Add auto_close_date support for equities --- tests/test_algorithm.py | 613 +++++++++++++++++- tests/test_assets.py | 54 +- tests/test_blotter.py | 34 + zipline/assets/asset_db_migrations.py | 28 + zipline/assets/asset_db_schema.py | 3 +- zipline/assets/asset_writer.py | 4 +- zipline/finance/blotter.py | 17 + .../finance/performance/position_tracker.py | 69 +- zipline/finance/performance/tracker.py | 24 +- zipline/finance/trading.py | 50 +- zipline/gens/tradesimulation.py | 135 +++- zipline/test_algorithms.py | 10 - zipline/utils/test_utils.py | 13 +- 13 files changed, 856 insertions(+), 198 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 5b0bc476..ec5f0dc5 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.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. +from collections import namedtuple import datetime from datetime import timedelta from mock import MagicMock @@ -22,13 +23,17 @@ from unittest import TestCase import numpy as np import pandas as pd +from contextlib2 import ExitStack +from zipline.api import FixedSlippage from zipline.assets import Equity, Future from zipline.utils.api_support import ZiplineAPI from zipline.utils.control_flow import nullctx from zipline.utils.test_utils import ( setup_logger, - teardown_logger + teardown_logger, + make_trade_panel_for_asset_info, + parameter_space, ) import zipline.utils.factory as factory import zipline.utils.simfactory as simfactory @@ -60,7 +65,6 @@ from zipline.test_algorithms import ( TestTargetAlgorithm, TestTargetPercentAlgorithm, TestTargetValueAlgorithm, - TestRemoveDataAlgo, SetLongOnlyAlgorithm, SetAssetDateBoundsAlgorithm, SetMaxPositionSizeAlgorithm, @@ -86,6 +90,8 @@ import zipline.utils.events from zipline.utils.test_utils import ( assert_single_position, drain_zipline, + make_jagged_equity_info, + tmp_asset_finder, to_utc, ) @@ -96,12 +102,14 @@ from zipline.sources import (SpecificEquityTrades, from zipline.finance.execution import LimitOrder from zipline.finance.trading import SimulationParameters +from zipline.finance.order import ORDER_STATUS from zipline.utils.api_support import set_algo_instance from zipline.utils.events import DateRuleFactory, TimeRuleFactory, Always from zipline.algorithm import TradingAlgorithm from zipline.protocol import DATASOURCE_TYPE from zipline.finance.trading import TradingEnvironment from zipline.finance.commission import PerShare +from zipline.utils.tradingcalendar import trading_day, trading_days # Because test cases appear to reuse some resources. _multiprocess_can_split_ = False @@ -1818,18 +1826,20 @@ class TestClosePosAlgo(TestCase): def setUp(self): self.env = TradingEnvironment() - self.days = self.env.trading_days[:4] + self.days = self.env.trading_days[:5] self.panel = pd.Panel({1: pd.DataFrame({ - 'price': [1, 1, 2, 4], 'volume': [1e9, 1e9, 1e9, 0], + 'price': [1, 1, 2, 4, 8], 'volume': [1e9, 1e9, 1e9, 1e9, 0], 'type': [DATASOURCE_TYPE.TRADE, + DATASOURCE_TYPE.TRADE, DATASOURCE_TYPE.TRADE, DATASOURCE_TYPE.TRADE, DATASOURCE_TYPE.CLOSE_POSITION]}, index=self.days) }) self.no_close_panel = pd.Panel({1: pd.DataFrame({ - 'price': [1, 1, 2, 4], 'volume': [1e9, 1e9, 1e9, 1e9], + 'price': [1, 1, 2, 4, 8], 'volume': [1e9, 1e9, 1e9, 1e9, 1e9], 'type': [DATASOURCE_TYPE.TRADE, + DATASOURCE_TYPE.TRADE, DATASOURCE_TYPE.TRADE, DATASOURCE_TYPE.TRADE, DATASOURCE_TYPE.TRADE]}, @@ -1838,7 +1848,7 @@ class TestClosePosAlgo(TestCase): def test_close_position_equity(self): metadata = {1: {'symbol': 'TEST', - 'end_date': self.days[3]}} + 'end_date': self.days[4]}} self.env.write_data(equities_data=metadata) algo = TestAlgorithm(sid=1, amount=1, order_count=1, commission=PerShare(0), @@ -1846,8 +1856,8 @@ class TestClosePosAlgo(TestCase): data = DataPanelSource(self.panel) # Check results - expected_positions = [0, 1, 1, 0] - expected_pnl = [0, 0, 1, 2] + expected_positions = [0, 1, 1, 1, 0] + expected_pnl = [0, 0, 1, 2, 4] results = algo.run(data) self.check_algo_positions(results, expected_positions) self.check_algo_pnl(results, expected_pnl) @@ -1861,8 +1871,8 @@ class TestClosePosAlgo(TestCase): data = DataPanelSource(self.panel) # Check results - expected_positions = [0, 1, 1, 0] - expected_pnl = [0, 0, 1, 2] + expected_positions = [0, 1, 1, 1, 0] + expected_pnl = [0, 0, 1, 2, 4] results = algo.run(data) self.check_algo_pnl(results, expected_pnl) self.check_algo_positions(results, expected_positions) @@ -1879,10 +1889,10 @@ class TestClosePosAlgo(TestCase): # Check results results = algo.run(data) - expected_positions = [0, 1, 1, 0] + expected_positions = [0, 1, 1, 1, 0] self.check_algo_positions(results, expected_positions) - expected_pnl = [0, 0, 1, 2] + expected_pnl = [0, 0, 1, 2, 0] self.check_algo_pnl(results, expected_pnl) def check_algo_pnl(self, results, expected_pnl): @@ -1982,39 +1992,578 @@ class TestTradingAlgorithm(TestCase): class TestRemoveData(TestCase): """ - tests if futures data is removed after expiry + tests if futures data is removed after max(expiration_date, end_date) """ def setUp(self): - dt = pd.Timestamp('2015-01-02', tz='UTC') - env = TradingEnvironment() - ix = env.trading_days.get_loc(dt) + self.env = env = TradingEnvironment() + start_date = pd.Timestamp('2015-01-02', tz='UTC') + start_ix = env.trading_days.get_loc(start_date) + days = env.trading_days - 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]}} + metadata = { + 0: { + 'symbol': 'X', + 'start_date': env.trading_days[start_ix + 2], + 'expiration_date': env.trading_days[start_ix + 5], + 'end_date': env.trading_days[start_ix + 6], + }, + 1: { + 'symbol': 'Y', + 'start_date': env.trading_days[start_ix + 4], + 'expiration_date': env.trading_days[start_ix + 7], + 'end_date': env.trading_days[start_ix + 8], + } + } env.write_data(futures_data=metadata) + assetX, assetY = env.asset_finder.retrieve_all([0, 1]) - index_x = env.trading_days[ix:ix + 5] + index_x = days[days.slice_indexer(assetX.start_date, assetX.end_date)] 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) + + index_y = days[days.slice_indexer(assetY.start_date, assetY.end_date)] 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) + self.trade_data = pd.Panel({0: data_x, 1: data_y}) + self.live_asset_counts = [] + assets = env.asset_finder.retrieve_all([0, 1]) + for day in self.trade_data.major_axis: + count = 0 + for asset in assets: + # We shouldn't see assets on their expiration dates. + if asset.start_date <= day <= asset.end_date: + count += 1 + self.live_asset_counts.append(count) def test_remove_data(self): - self.algo.run(self.source) + source = DataPanelSource(self.trade_data) - 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) + def initialize(context): + context.data_lengths = [] + + def handle_data(context, data): + context.data_lengths.append(len(data)) + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + env=self.env, + ) + + algo.run(source) + self.assertEqual(algo.data_lengths, self.live_asset_counts) + + +class TestEquityAutoClose(TestCase): + """ + Tests if delisted equities are properly removed from a portfolio holding + positions in said equities. + """ + @classmethod + def setUpClass(cls): + start_date = pd.Timestamp('2015-01-05', tz='UTC') + start_date_loc = trading_days.get_loc(start_date) + test_duration = 7 + cls.test_days = trading_days[ + start_date_loc:start_date_loc + test_duration + ] + cls.first_asset_expiration = cls.test_days[2] + + def setUp(self): + self._teardown_stack = ExitStack() + + def tearDown(self): + self._teardown_stack.close() + + def make_temp_resource(self, resource_context): + return self._teardown_stack.enter_context(resource_context) + + def make_data(self, auto_close_delta, frequency): + asset_info = make_jagged_equity_info( + num_assets=3, + start_date=self.test_days[0], + first_end=self.first_asset_expiration, + frequency=trading_day, + periods_between_ends=2, + auto_close_delta=auto_close_delta, + ) + + # Manually set the trading environment's asset finder. + finder = self.make_temp_resource(tmp_asset_finder(equities=asset_info)) + sids = list(asset_info.index) + assets = finder.retrieve_all(sids) + env = TradingEnvironment(asset_db_path=None) + env.asset_finder = finder + + if frequency == 'daily': + dates = self.test_days + elif frequency == 'minute': + dates = env.minutes_for_days_in_range( + self.test_days[0], + self.test_days[-1], + ) + else: + self.fail("Unknown frequency in make_data: %r" % frequency) + + prices_and_volumes = make_trade_panel_for_asset_info( + dates=dates, + asset_info=asset_info, + price_start=10, + price_step_by_sid=10, + price_step_by_date=1, + volume_start=100, + volume_step_by_sid=100, + volume_step_by_date=10, + ) + + if frequency == 'daily': + final_prices = { + asset.sid: prices_and_volumes.loc[ + asset.sid, + asset.end_date, + 'price', + ] + for asset in assets + } + else: + final_prices = { + asset.sid: prices_and_volumes.loc[ + asset.sid, + env.get_open_and_close(asset.end_date)[1], + 'price', + ] + for asset in assets + } + + TestData = namedtuple( + 'TestData', + [ + 'asset_info', + 'assets', + 'env', + 'final_prices', + 'finder', + 'prices_and_volumes', + ], + ) + return TestData( + asset_info=asset_info, + assets=assets, + env=env, + final_prices=final_prices, + finder=finder, + prices_and_volumes=prices_and_volumes, + ) + + def prices_on_tick(self, prices_and_volumes, N): + return prices_and_volumes.ix[ + :, N, 'price' + ] + + def default_initialize(self): + """ + Initialize function shared between test algos. + """ + def initialize(context): + context.ordered = False + context.set_commission(PerShare(0)) + context.set_slippage(FixedSlippage(spread=0)) + context.num_positions = [] + context.cash = [] + + return initialize + + def default_handle_data(self, assets, order_size): + """ + Handle data function shared between test algos. + """ + def handle_data(context, data): + if not context.ordered: + for asset in assets: + context.order(asset, order_size) + context.ordered = True + + context.cash.append(context.portfolio.cash) + context.num_positions.append(len(context.portfolio.positions)) + + return handle_data + + @parameter_space( + order_size=[10, -10], + capital_base=[0, 100000], + auto_close_lag=[1, 2], + ) + def test_daily_delisted_equities(self, + order_size, + capital_base, + auto_close_lag): + """ + Make sure that after an equity gets delisted, our portfolio holds the + correct number of equities and correct amount of cash. + """ + auto_close_delta = trading_day * auto_close_lag + resources = self.make_data(auto_close_delta, 'daily') + + assets = resources.assets + sids = [asset.sid for asset in assets] + env = resources.env + prices_and_volumes = resources.prices_and_volumes + final_prices = resources.final_prices + + source = DataPanelSource(prices_and_volumes) + + # Prices at which we expect our orders to be filled. + initial_fill_prices = self.prices_on_tick(prices_and_volumes, 1) + cost_basis = sum(initial_fill_prices) * order_size + + # Last known prices of assets that will be auto-closed. + fp0 = final_prices[0] + fp1 = final_prices[1] + + algo = TradingAlgorithm( + initialize=self.default_initialize(), + handle_data=self.default_handle_data(assets, order_size), + env=env, + capital_base=capital_base, + ) + output = algo.run(source) + + initial_cash = capital_base + after_fills = initial_cash - cost_basis + after_first_auto_close = after_fills + fp0 * (order_size) + after_second_auto_close = after_first_auto_close + fp1 * (order_size) + + if auto_close_lag == 1: + # Day 1: Order 10 shares of each equity; there are 3 equities. + # Day 2: Order goes through at the day 2 price of each equity. + # Day 3: End date of Equity 0. + # Day 4: Auto close date of Equity 0. Add cash == (fp0 * size). + # Day 5: End date of Equity 1. + # Day 6: Auto close date of Equity 1. Add cash == (fp1 * size). + # Day 7: End date of Equity 2 and last day of backtest; no changes. + expected_cash = [ + initial_cash, + after_fills, + after_fills, + after_first_auto_close, + after_first_auto_close, + after_second_auto_close, + after_second_auto_close, + ] + expected_num_positions = [0, 3, 3, 2, 2, 1, 1] + elif auto_close_lag == 2: + # Day 1: Order 10 shares of each equity; there are 3 equities. + # Day 2: Order goes through at the day 2 price of each equity. + # Day 3: End date of Equity 0. + # Day 4: Nothing happens. + # Day 5: End date of Equity 1. Auto close of equity 0. + # Add cash == (fp0 * size). + # Day 6: Nothing happens. + # Day 7: End date of Equity 2 and auto-close date of Equity 1. + # Add cash equal to (fp1 * size). + expected_cash = [ + initial_cash, + after_fills, + after_fills, + after_fills, + after_first_auto_close, + after_first_auto_close, + after_second_auto_close, + ] + expected_num_positions = [0, 3, 3, 3, 2, 2, 1] + else: + self.fail( + "Don't know about auto_close lags other than 1 or 2. " + "Add test answers please!" + ) + + # Check expected cash. + self.assertEqual(algo.cash, expected_cash) + self.assertEqual(expected_cash, list(output['ending_cash'])) + + # Check expected long/short counts. + # We have longs if order_size > 0. + # We have shrots if order_size < 0. + self.assertEqual(algo.num_positions, expected_num_positions) + if order_size > 0: + self.assertEqual( + expected_num_positions, + list(output['longs_count']), + ) + self.assertEqual( + [0] * len(self.test_days), + list(output['shorts_count']), + ) + else: + self.assertEqual( + expected_num_positions, + list(output['shorts_count']), + ) + self.assertEqual( + [0] * len(self.test_days), + list(output['longs_count']), + ) + + # Check expected transactions. + # We should have a transaction of order_size shares per sid. + transactions = output['transactions'] + initial_fills = transactions.iloc[1] + self.assertEqual(len(initial_fills), len(assets)) + for sid, txn in zip(sids, initial_fills): + self.assertDictContainsSubset( + { + 'amount': order_size, + 'commission': 0.0, + 'dt': self.test_days[1], + 'price': initial_fill_prices[sid], + 'sid': sid, + }, + txn, + ) + # This will be a UUID. + self.assertIsInstance(txn['order_id'], str) + + def transactions_for_date(date): + return transactions.iloc[self.test_days.get_loc(date)] + + # We should have exactly one auto-close transaction on the close date + # of asset 0. + (first_auto_close_transaction,) = transactions_for_date( + assets[0].auto_close_date + ) + self.assertEqual( + first_auto_close_transaction, + { + 'amount': -order_size, + 'commission': 0.0, + 'dt': assets[0].auto_close_date, + 'price': fp0, + 'sid': sids[0], + 'order_id': None, # Auto-close txns emit Nones for order_id. + }, + ) + + (second_auto_close_transaction,) = transactions_for_date( + assets[1].auto_close_date + ) + self.assertEqual( + second_auto_close_transaction, + { + 'amount': -order_size, + 'commission': 0.0, + 'dt': assets[1].auto_close_date, + 'price': fp1, + 'sid': sids[1], + 'order_id': None, # Auto-close txns emit Nones for order_id. + }, + ) + + def test_cancel_open_orders(self): + """ + Test that any open orders for an equity that gets delisted are + canceled. Unless an equity is auto closed, any open orders for that + equity will persist indefinitely. + """ + auto_close_delta = trading_day + resources = self.make_data(auto_close_delta, 'daily') + env = resources.env + assets = resources.assets + + source = DataPanelSource(resources.prices_and_volumes) + + first_asset_end_date = assets[0].end_date + first_asset_auto_close_date = assets[0].auto_close_date + + def initialize(context): + pass + + def handle_data(context, data): + # The only order we place in this test should never be filled. + assert ( + context.portfolio.cash == context.portfolio.starting_cash + ) + + now = context.get_datetime() + + if now == first_asset_end_date: + # Equity 0 will no longer exist tomorrow, so this order will + # never be filled. + assert len(context.get_open_orders()) == 0 + context.order(context.sid(0), 10) + assert len(context.get_open_orders()) == 1 + elif now == first_asset_auto_close_date: + assert len(context.get_open_orders()) == 0 + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + env=env, + ) + results = algo.run(source) + + orders = results['orders'] + + def orders_for_date(date): + return orders.iloc[self.test_days.get_loc(date)] + + original_open_orders = orders_for_date(first_asset_end_date) + assert len(original_open_orders) == 1 + self.assertDictContainsSubset( + { + 'amount': 10, + 'commission': None, + 'created': first_asset_end_date, + 'dt': first_asset_end_date, + 'sid': assets[0], + 'status': ORDER_STATUS.OPEN, + 'filled': 0, + }, + original_open_orders[0], + ) + + orders_after_auto_close = orders_for_date(first_asset_auto_close_date) + assert len(orders_after_auto_close) == 1 + self.assertDictContainsSubset( + { + 'amount': 10, + 'commission': None, + 'created': first_asset_end_date, + 'dt': first_asset_auto_close_date, + 'sid': assets[0], + 'status': ORDER_STATUS.CANCELLED, + 'filled': 0, + }, + orders_after_auto_close[0], + ) + + def test_minutely_delisted_equities(self): + resources = self.make_data(trading_day, 'minute') + + env = resources.env + assets = resources.assets + sids = [a.sid for a in assets] + final_prices = resources.final_prices + prices_and_volumes = resources.prices_and_volumes + backtest_minutes = prices_and_volumes.major_axis + + order_size = 10 + source = DataPanelSource(prices_and_volumes) + + capital_base = 100000 + algo = TradingAlgorithm( + initialize=self.default_initialize(), + handle_data=self.default_handle_data(assets, order_size), + env=env, + data_frequency='minute', + capital_base=capital_base, + ) + output = algo.run(source) + initial_fill_prices = self.prices_on_tick(prices_and_volumes, 1) + cost_basis = sum(initial_fill_prices) * order_size + + # Last known prices of assets that will be auto-closed. + fp0 = final_prices[0] + fp1 = final_prices[1] + + initial_cash = capital_base + after_fills = initial_cash - cost_basis + after_first_auto_close = after_fills + fp0 * (order_size) + after_second_auto_close = after_first_auto_close + fp1 * (order_size) + + expected_cash = [initial_cash] + expected_position_counts = [0] + + # We have the rest of the first sim day, plus the second and third + # days' worth of minutes with cash spent. + expected_cash.extend([after_fills] * (389 + 390 + 390)) + expected_position_counts.extend([3] * (389 + 390 + 390)) + + # We then have two days with the cash refunded from asset 0. + expected_cash.extend([after_first_auto_close] * (390 + 390)) + expected_position_counts.extend([2] * (390 + 390)) + + # We then have two days with cash refunded from asset 1 + expected_cash.extend([after_second_auto_close] * (390 + 390)) + expected_position_counts.extend([1] * (390 + 390)) + + self.assertEqual(algo.cash, expected_cash) + self.assertEqual( + list(output['ending_cash']), + [ + after_fills, + after_fills, + after_fills, + after_first_auto_close, + after_first_auto_close, + after_second_auto_close, + after_second_auto_close, + ], + ) + + self.assertEqual(algo.num_positions, expected_position_counts) + self.assertEqual( + list(output['longs_count']), + [3, 3, 3, 2, 2, 1, 1], + ) + + # Check expected transactions. + # We should have a transaction of order_size shares per sid. + transactions = output['transactions'] + + # Note that the transactions appear on the first day rather than the + # second in minute mode, because the fills happen on the second tick of + # the backtest, which is still on the first day in minute mode. + initial_fills = transactions.iloc[0] + self.assertEqual(len(initial_fills), len(assets)) + for sid, txn in zip(sids, initial_fills): + self.assertDictContainsSubset( + { + 'amount': order_size, + 'commission': 0.0, + 'dt': backtest_minutes[1], + 'price': initial_fill_prices[sid], + 'sid': sid, + }, + txn, + ) + # This will be a UUID. + self.assertIsInstance(txn['order_id'], str) + + def transactions_for_date(date): + return transactions.iloc[self.test_days.get_loc(date)] + + # We should have exactly one auto-close transaction on the close date + # of asset 0. + (first_auto_close_transaction,) = transactions_for_date( + assets[0].auto_close_date + ) + self.assertEqual( + first_auto_close_transaction, + { + 'amount': -order_size, + 'commission': 0.0, + 'dt': assets[0].auto_close_date, + 'price': fp0, + 'sid': sids[0], + 'order_id': None, # Auto-close txns emit Nones for order_id. + }, + ) + + (second_auto_close_transaction,) = transactions_for_date( + assets[1].auto_close_date + ) + self.assertEqual( + second_auto_close_transaction, + { + 'amount': -order_size, + 'commission': 0.0, + 'dt': assets[1].auto_close_date, + 'price': fp1, + 'sid': sids[1], + 'order_id': None, # Auto-close txns emit Nones for order_id. + }, + ) diff --git a/tests/test_assets.py b/tests/test_assets.py index 0a90664e..ff3337c7 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -20,6 +20,7 @@ from contextlib import contextmanager from datetime import datetime, timedelta import pickle import sys +from types import GetSetDescriptorType from unittest import TestCase import uuid import warnings @@ -172,6 +173,22 @@ def build_lookup_generic_cases(asset_finder_type): class AssetTestCase(TestCase): + # Dynamically list the Asset properties we want to test. + asset_attrs = [name for name, value in vars(Asset).items() + if isinstance(value, GetSetDescriptorType)] + + # Very wow + asset = Asset( + 1337, + symbol="DOGE", + asset_name="DOGECOIN", + start_date=pd.Timestamp('2013-12-08 9:31AM', tz='UTC'), + end_date=pd.Timestamp('2014-06-25 11:21AM', tz='UTC'), + first_traded=pd.Timestamp('2013-12-08 9:31AM', tz='UTC'), + auto_close_date=pd.Timestamp('2014-06-26 11:21AM', tz='UTC'), + exchange='THE MOON', + ) + def test_asset_object(self): self.assertEquals({5061: 'foo'}[Asset(5061)], 'foo') self.assertEquals(Asset(5061), 5061) @@ -182,32 +199,19 @@ class AssetTestCase(TestCase): self.assertEquals(str(Asset(5061)), 'Asset(5061)') + def test_to_and_from_dict(self): + asset_from_dict = Asset.from_dict(self.asset.to_dict()) + for attr in self.asset_attrs: + self.assertEqual( + getattr(self.asset, attr), getattr(asset_from_dict, attr), + ) + def test_asset_is_pickleable(self): - - # Very wow - s = Asset( - 1337, - symbol="DOGE", - asset_name="DOGECOIN", - start_date=pd.Timestamp('2013-12-08 9:31AM', tz='UTC'), - end_date=pd.Timestamp('2014-06-25 11:21AM', tz='UTC'), - first_traded=pd.Timestamp('2013-12-08 9:31AM', tz='UTC'), - exchange='THE MOON', - ) - s_unpickled = pickle.loads(pickle.dumps(s)) - - attrs_to_check = ['end_date', - 'exchange', - 'first_traded', - 'end_date', - 'asset_name', - 'start_date', - 'sid', - 'start_date', - 'symbol'] - - for attr in attrs_to_check: - self.assertEqual(getattr(s, attr), getattr(s_unpickled, attr)) + asset_unpickled = pickle.loads(pickle.dumps(self.asset)) + for attr in self.asset_attrs: + self.assertEqual( + getattr(self.asset, attr), getattr(asset_unpickled, attr), + ) def test_asset_comparisons(self): diff --git a/tests/test_blotter.py b/tests/test_blotter.py index 6f2ffee1..f9606b09 100644 --- a/tests/test_blotter.py +++ b/tests/test_blotter.py @@ -65,6 +65,40 @@ class BlotterTestCase(TestCase): self.assertEqual(result.limit, expected_lmt) self.assertEqual(result.stop, expected_stp) + def test_cancel(self): + blotter = Blotter() + + oid_1 = blotter.order(24, 100, MarketOrder()) + oid_2 = blotter.order(24, 200, MarketOrder()) + oid_3 = blotter.order(24, 300, MarketOrder()) + + # Create an order for another asset to verify that we don't remove it + # when we do cancel_all on 24. + blotter.order(25, 150, MarketOrder()) + + self.assertEqual(len(blotter.open_orders), 2) + self.assertEqual(len(blotter.open_orders[24]), 3) + self.assertEqual( + [o.amount for o in blotter.open_orders[24]], + [100, 200, 300], + ) + + blotter.cancel(oid_2) + self.assertEqual(len(blotter.open_orders), 2) + self.assertEqual(len(blotter.open_orders[24]), 2) + self.assertEqual( + [o.amount for o in blotter.open_orders[24]], + [100, 300], + ) + self.assertEqual( + [o.id for o in blotter.open_orders[24]], + [oid_1, oid_3], + ) + + blotter.cancel_all(24) + self.assertEqual(len(blotter.open_orders), 1) + self.assertEqual(list(blotter.open_orders), [25]) + def test_order_rejection(self): blotter = Blotter() # Reject a nonexistent order -> no order appears in new_order, diff --git a/zipline/assets/asset_db_migrations.py b/zipline/assets/asset_db_migrations.py index b1c5aa2f..18ed61d2 100644 --- a/zipline/assets/asset_db_migrations.py +++ b/zipline/assets/asset_db_migrations.py @@ -101,9 +101,37 @@ def _downgrade_v1_to_v0(op, version_info_table): write_version_info(version_info_table, 0) + +def _downgrade_v2_to_v1(op, version_info_table): + """ + Downgrade assets db by removing the 'auto_close_date' column. + """ + version_info_table.delete().execute() + + # Drop indices before batch + # This is to prevent index collision when creating the temp table + op.drop_index('ix_equities_fuzzy_symbol') + op.drop_index('ix_equities_company_symbol') + + # Execute batch op to allow column modification in SQLite + with op.batch_alter_table('equities') as batch_op: + + batch_op.drop_column('auto_close_date') + + # Recreate indices after batch + op.create_index('ix_equities_fuzzy_symbol', + table_name='equities', + columns=['fuzzy_symbol']) + op.create_index('ix_equities_company_symbol', + table_name='equities', + columns=['company_symbol']) + + write_version_info(version_info_table, 1) + # This dict contains references to downgrade methods that can be applied to an # assets db. The resulting db's version is the key. # e.g. The method at key '0' is the downgrade method from v1 to v0 _downgrade_methods = { 0: _downgrade_v1_to_v0, + 1: _downgrade_v2_to_v1, } diff --git a/zipline/assets/asset_db_schema.py b/zipline/assets/asset_db_schema.py index 3e0b4f15..bc79f5e7 100644 --- a/zipline/assets/asset_db_schema.py +++ b/zipline/assets/asset_db_schema.py @@ -4,7 +4,7 @@ import sqlalchemy as sa # Define a version number for the database generated by these writers # Increment this version number any time a change is made to the schema of the # assets database -ASSET_DB_VERSION = 1 +ASSET_DB_VERSION = 2 def generate_asset_db_metadata(bind=None): @@ -46,6 +46,7 @@ def _equities_table_schema(metadata): sa.Column('start_date', sa.Integer, default=0, nullable=False), sa.Column('end_date', sa.Integer, nullable=False), sa.Column('first_traded', sa.Integer, nullable=False), + sa.Column('auto_close_date', sa.Integer), sa.Column('exchange', sa.Text), ) diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 96d1c022..76b82293 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -45,6 +45,7 @@ _equities_defaults = { 'start_date': 0, 'end_date': 2 ** 62 - 1, 'first_traded': None, + 'auto_close_date': None, 'exchange': None, } @@ -403,7 +404,8 @@ class AssetDBWriter(with_metaclass(ABCMeta)): equities_output.fuzzy_symbol.str.upper() # Convert date columns to UNIX Epoch integers (nanoseconds) - for date_col in ('start_date', 'end_date', 'first_traded'): + for date_col in ('start_date', 'end_date', 'first_traded', + 'auto_close_date'): equities_output[date_col] = \ self.dt_to_epoch_ns(equities_output[date_col]) diff --git a/zipline/finance/blotter.py b/zipline/finance/blotter.py index 320fb12a..26fd133a 100644 --- a/zipline/finance/blotter.py +++ b/zipline/finance/blotter.py @@ -127,6 +127,23 @@ class Blotter(object): # along with newly placed orders. self.new_orders.append(cur_order) + def cancel_all(self, sid): + """ + Cancel all open orders for a given sid. + """ + # (sadly) open_orders is a defaultdict, so this will always succeed. + orders = self.open_orders[sid] + + # We're making a copy here because `cancel` mutates the list of open + # orders in place. The right thing to do here would be to make + # self.open_orders no longer a defaultdict. If we do that, then we + # should just remove the orders once here and be done with the matter. + for order in orders[:]: + self.cancel(order.id) + + assert not orders + del self.open_orders[sid] + def reject(self, order_id, reason=''): """ Mark the given order as 'rejected', which is functionally similar to diff --git a/zipline/finance/performance/position_tracker.py b/zipline/finance/performance/position_tracker.py index 8489ccb6..1697002f 100644 --- a/zipline/finance/performance/position_tracker.py +++ b/zipline/finance/performance/position_tracker.py @@ -27,7 +27,6 @@ except ImportError: from collections import OrderedDict from six import iteritems, itervalues -from zipline.protocol import Event, DATASOURCE_TYPE from zipline.finance.transaction import Transaction from zipline.utils.serialization_utils import ( VERSION_LABEL @@ -136,10 +135,6 @@ class PositionTracker(object): ) self._positions_store = zp.Positions() - # Dict, keyed on dates, that contains lists of close position events - # for any Assets in this tracker's positions - self._auto_close_position_sids = {} - def _update_asset(self, sid): try: self._position_value_multipliers[sid] @@ -157,64 +152,6 @@ class PositionTracker(object): if isinstance(asset, Future): self._position_value_multipliers[sid] = 0 self._position_exposure_multipliers[sid] = asset.multiplier - # Futures auto-close timing is controlled by the Future's - # auto_close_date property - self._insert_auto_close_position_date( - dt=asset.auto_close_date, - sid=sid - ) - - def _insert_auto_close_position_date(self, dt, sid): - """ - Inserts the given SID in to the list of positions to be auto-closed by - the given dt. - - Parameters - ---------- - dt : pandas.Timestamp - The date before-which the given SID will be auto-closed - sid : int - The SID of the Asset to be auto-closed - """ - if dt is not None: - self._auto_close_position_sids.setdefault(dt, set()).add(sid) - - def auto_close_position_events(self, next_trading_day): - """ - Generates CLOSE_POSITION events for any SIDs whose auto-close date is - before or equal to the given date. - - Parameters - ---------- - next_trading_day : pandas.Timestamp - The time before-which certain Assets need to be closed - - Yields - ------ - Event - A close position event for any sids that should be closed before - the next_trading_day parameter - """ - past_asset_end_dates = set() - - # Check the auto_close_position_dates dict for SIDs to close - for date, sids in self._auto_close_position_sids.items(): - if date > next_trading_day: - continue - past_asset_end_dates.add(date) - - for sid in sids: - # Yield a CLOSE_POSITION event - event = Event({ - 'dt': date, - 'type': DATASOURCE_TYPE.CLOSE_POSITION, - 'sid': sid, - }) - yield event - - # Clear out past dates - while past_asset_end_dates: - self._auto_close_position_sids.pop(past_asset_end_dates.pop()) def update_last_sale(self, event): # NOTE, PerformanceTracker already vetted as TRADE type @@ -362,7 +299,7 @@ class PositionTracker(object): dt=event.dt, price=price, commission=0, - order_id=0 + order_id=None, ) return txn @@ -447,9 +384,8 @@ class PositionTracker(object): state_dict['asset_finder'] = self.asset_finder state_dict['positions'] = dict(self.positions) state_dict['unpaid_dividends'] = self._unpaid_dividends - state_dict['auto_close_position_sids'] = self._auto_close_position_sids - STATE_VERSION = 3 + STATE_VERSION = 4 state_dict[VERSION_LABEL] = STATE_VERSION return state_dict @@ -467,7 +403,6 @@ class PositionTracker(object): self._positions_store = zp.Positions() self._unpaid_dividends = state['unpaid_dividends'] - self._auto_close_position_sids = state['auto_close_position_sids'] # Arrays for quick calculations of positions value self._position_value_multipliers = OrderedDict() diff --git a/zipline/finance/performance/tracker.py b/zipline/finance/performance/tracker.py index a084bdfc..0b0182de 100644 --- a/zipline/finance/performance/tracker.py +++ b/zipline/finance/performance/tracker.py @@ -391,22 +391,6 @@ class PerformanceTracker(object): self.cumulative_performance.handle_dividends_paid(net_cash_payment) self.todays_performance.handle_dividends_paid(net_cash_payment) - def check_asset_auto_closes(self, next_trading_day): - """ - Check if the position tracker currently owns any Assets with an - auto-close date that is the next trading day. Close those positions. - - Parameters - ---------- - next_trading_day : pandas.Timestamp - The next trading day of the simulation - """ - auto_close_events = self.position_tracker.auto_close_position_events( - next_trading_day=next_trading_day - ) - for event in auto_close_events: - self.process_close_position(event) - def handle_minute_close(self, dt): """ Handles the close of the given minute. This includes handling @@ -473,11 +457,6 @@ class PerformanceTracker(object): # simulation, return the daily perf packet next_trading_day = self.env.next_trading_day(completed_date) - # Check if any assets need to be auto-closed before generating today's - # perf period - if next_trading_day: - self.check_asset_auto_closes(next_trading_day=next_trading_day) - # Take a snapshot of our current performance to return to the # browser. daily_update = self.to_dict(emission_type='daily') @@ -502,8 +481,7 @@ class PerformanceTracker(object): if (next_trading_day is None) or (next_trading_day >= self.last_close): return daily_update - # Check for any dividends and auto-closes, then return the daily perf - # packet + # Check for any dividends, then return the daily perf packet self.check_upcoming_dividends(next_trading_day=next_trading_day) return daily_update diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index dd436fdc..d914c611 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -174,23 +174,43 @@ class TradingEnvironment(object): # If any pandas.DataFrame data has been provided, # write it to the database. - if (equities_df is not None or futures_df is not None or - exchanges_df is not None or root_symbols_df is not None): - self._write_data_dataframes(equities_df, futures_df, - exchanges_df, root_symbols_df) + has_rows = lambda df: df is not None and len(df) > 0 + if any(map(has_rows, [equities_df, + futures_df, + exchanges_df, + root_symbols_df])): + self._write_data_dataframes( + equities=equities_df, + futures=futures_df, + exchanges=exchanges_df, + root_symbols=root_symbols_df, + ) - if (equities_data is not None or futures_data is not None or - exchanges_data is not None or root_symbols_data is not None): - self._write_data_dicts(equities_data, futures_data, - exchanges_data, root_symbols_data) + # Same for dicts. + has_data = lambda d: d is not None and len(d) > 0 + if any(map(has_data, [equities_data, + futures_data, + exchanges_data, + futures_data])): + self._write_data_dicts( + equities=equities_data, + futures=futures_data, + exchanges=exchanges_data, + root_symbols=root_symbols_data + ) - # These could be lists or other iterables such as a pandas.Index. - # For simplicity, don't check whether data has been provided. - self._write_data_lists(equities_identifiers, - futures_identifiers, - exchanges_identifiers, - root_symbols_identifiers, - allow_sid_assignment=allow_sid_assignment) + # Same for iterables. + if any(map(has_data, [equities_identifiers, + futures_identifiers, + exchanges_identifiers, + root_symbols_identifiers])): + self._write_data_lists( + equities=equities_identifiers, + futures=futures_identifiers, + exchanges=exchanges_identifiers, + root_symbols=root_symbols_identifiers, + allow_sid_assignment=allow_sid_assignment + ) def _write_data_lists(self, equities=None, futures=None, exchanges=None, root_symbols=None, allow_sid_assignment=True): diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 23ae30db..e6dc8c54 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -23,8 +23,9 @@ from zipline.errors import SidsNotFound from zipline.finance.trading import NoFurtherDataError from zipline.protocol import ( BarData, + DATASOURCE_TYPE, + Event, SIDData, - DATASOURCE_TYPE ) from zipline.utils.api_support import ZiplineAPI from zipline.utils.data import SortedDict @@ -58,10 +59,35 @@ class AlgorithmSimulator(object): # Snapshot Setup # ============== - def _get_asset_close_date(sid, - finder=self.env.asset_finder, - default=self.sim_params.last_close - + timedelta(days=1)): + _day = timedelta(days=1) + + def _get_removal_date(sid, + finder=self.env.asset_finder, + default=self.sim_params.last_close + _day): + """ + Get the date of the morning on which we should remove an asset from + data. + + If we don't have an auto_close_date, this is just the end of the + simulation. + + If we have an auto_close_date, then we remove assets from data on + max(asset.auto_close_date, asset.end_date + timedelta(days=1)) + + We hold assets at least until auto_close_date because up until that + date the user might still hold positions or have open orders in an + expired asset. + + We hold assets at least until end_date + 1, because an asset + continues trading until the **end** of its end_date. Even if an + asset auto-closed before the end_date (say, because Interactive + Brokers clears futures positions prior the actual notice or + expiration), there may still be trades arriving that represent + signals for other assets that are still tradeable. (Particularly in + the futures case, trading in the final days of a contract are + likely relevant for trading the next contract on the same future + chain.) + """ try: asset = finder.retrieve_asset(sid) except ValueError: @@ -72,18 +98,30 @@ class AlgorithmSimulator(object): return default + timedelta(microseconds=id(sid)) except SidsNotFound: return default - # Default is used when the asset has no auto close date, - # and is set to a time after the simulation ends, so that the - # relevant asset isn't removed from the universe at all - # (at least not for this reason). - return asset.auto_close_date or default - self._get_asset_close = _get_asset_close_date + auto_close_date = asset.auto_close_date + if auto_close_date is None: + # If we don't have an auto_close_date, we never remove an asset + # from the user's portfolio. + return default + + end_date = asset.end_date + if end_date is None: + # If we have an auto_close_date but not an end_date, clear the + # asset from data when we clear positions/orders. + return auto_close_date + + # If we have both, make close once we're on or after the + # auto_close_date, and strictly after the end_date. + # See docstring above for an explanation of this logic. + return max(auto_close_date, end_date + _day) + + self._get_removal_date = _get_removal_date # The algorithm's data as of our most recent event. # Maintain sids in order by asset close date, so that we can more # efficiently remove them when their times come... - self.current_data = BarData(SortedDict(self._get_asset_close)) + self.current_data = BarData(SortedDict(self._get_removal_date)) # We don't have a datetime for the current snapshot until we # receive a message. @@ -124,16 +162,6 @@ class AlgorithmSimulator(object): self.simulation_dt = date self.on_dt_changed(date) - closed = list(takewhile( - lambda asset_id: self._get_asset_close(asset_id) < date, - self.current_data - )) - for sid in closed: - 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. @@ -368,8 +396,71 @@ class AlgorithmSimulator(object): dt = normalize_date(dt) self.simulation_dt = dt self.on_dt_changed(dt) + + self._cleanup_expired_assets(dt, self.current_data, self.algo.blotter) + self.algo.before_trading_start(self.current_data) + def _cleanup_expired_assets(self, dt, current_data, algo_blotter): + """ + Clear out any assets that have expired before starting a new sim day. + + Performs three functions: + + 1. Finds all assets for which we have open orders and clears any + orders whose assets are on or after their auto_close_date. + + 2. Finds all assets for which we have positions and generates + close_position events for any assets that have reached their + auto_close_date. + + 3. Finds and removes from data all sids for which + _get_removal_date(sid) <= dt. + """ + algo = self.algo + expired = list(takewhile( + lambda asset_id: self._get_removal_date(asset_id) <= dt, + self.current_data + )) + for sid in expired: + try: + del self.current_data[sid] + except KeyError: + continue + + def create_close_position_event(asset): + event = Event({ + 'dt': dt, + 'type': DATASOURCE_TYPE.CLOSE_POSITION, + 'sid': asset.sid, + }) + return event + + def past_auto_close_date(asset): + acd = asset.auto_close_date + return acd is not None and acd <= dt + + # Remove positions in any sids that have reached their auto_close date. + to_clear = [] + finder = algo.asset_finder + position_assets = finder.retrieve_all(list(algo.portfolio.positions)) + for asset in position_assets: + if past_auto_close_date(asset): + to_clear.append(asset) + perf_tracker = algo.perf_tracker + for close_event in map(create_close_position_event, to_clear): + perf_tracker.process_close_position(close_event) + + # Remove open orders for any sids that have reached their + # auto_close_date. + blotter = algo.blotter + to_cancel = [] + for asset in blotter.open_orders: + if past_auto_close_date(asset): + to_cancel.append(asset) + for asset in to_cancel: + blotter.cancel_all(asset) + def on_dt_changed(self, dt): if self.algo.datetime != dt: self.algo.on_dt_changed(dt) diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 6d8d399b..04defd16 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -937,16 +937,6 @@ 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 diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py index 45e02cbb..0d51a6ca 100644 --- a/zipline/utils/test_utils.py +++ b/zipline/utils/test_utils.py @@ -15,6 +15,7 @@ import tempfile from logbook import FileHandler from mock import patch from numpy.testing import assert_allclose, assert_array_equal +import numpy as np import pandas as pd from pandas.tseries.offsets import MonthBegin from six import iteritems, itervalues @@ -351,7 +352,8 @@ def make_jagged_equity_info(num_assets, start_date, first_end, frequency, - periods_between_ends): + periods_between_ends, + auto_close_delta): """ Create a DataFrame representing assets that all begin at the same start date, but have cascading end dates. @@ -375,7 +377,7 @@ def make_jagged_equity_info(num_assets, info : pd.DataFrame DataFrame representing newly-created assets. """ - return pd.DataFrame( + frame = pd.DataFrame( { 'symbol': [chr(ord('A') + i) for i in range(num_assets)], 'start_date': start_date, @@ -389,6 +391,13 @@ def make_jagged_equity_info(num_assets, index=range(num_assets), ) + # Explicitly pass None to disable setting the auto_close_date column. + if auto_close_delta is not None: + frame['auto_close_date'] = frame['end_date'] + auto_close_delta + + return frame + + def make_trade_panel_for_asset_info(dates, asset_info, price_start,