wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-19 11:22:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

ENH: Add auto_close_date support for equities

Browse Source
This commit is contained in:
dmichalowicz
2016-02-22 13:51:20 -05:00
parent e2db1a7b45
commit 5be63f36d5
13 changed files with 856 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_algorithm.py
+581 -32
View File
@@ -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.
},
)