import numpy as np
def initialize(context):
# trading Coke (KO) and Pepsi (PEP)
context.stocks = [sid(4283), sid(5885)]
# $200,000 bets
context.bet_amount = 200000
# long (short) if KO/PEP is below (above) 20-day average +(-) 2 * 20-day standard deviation
# close long (short) position when KO/PEP rises (falls) to long (short) price -(+) 0.5 * long (short) standard deviation
context.trade_tolerance = 2.0
context.close_tolerance = 0.5
context.close_trades = {"long": {}, "short": {}}
def handle_data(context, data):
rval = avgstd(data,context)
if rval is None:
return
avg_spread, std_spread = rval
cur_spread = data[context.stocks[0]].price/data[context.stocks[1]].price
signal = calculate_direction(context, cur_spread, avg_spread, std_spread)
a_price, b_price = data[context.stocks[0]].price, data[context.stocks[1]].price
a_order, b_order = calculate_order(context, signal, a_price, b_price)
if a_order == 0 and b_order == 0: return
calculate_close_trades(context,signal,cur_spread,std_spread,a_order,b_order)
a_order, b_order = calculate_net_orders(context,cur_spread,a_order,b_order)
order(context.stocks[0], a_order)
order(context.stocks[1], b_order)
if signal:
logmsg = '\n{sig} a_price: {ap} b_price: {bp} cur_spread: {c} avg_spread: {m} std_spread: {s} a_position: {a} b_position: {b}\n'
log.info(logmsg.format(
sig=signal,
ap=a_price,
bp=b_price,
c=cur_spread,
m=avg_spread,
s=std_spread,
a=a_order,
b=b_order
))
def calculate_direction(context, cur_spread, avg_spread, std_spread):
""" calculates the direction of trade given trade tolerance, current spraed, average spread and standard deviation of spread """
if avg_spread is not None and (context.portfolio.positions_value + context.portfolio.cash) > 0:
if cur_spread >= avg_spread + std_spread * context.trade_tolerance:
# if 0 or negative holdings in A, short spread -- short A, long B
return "short"
elif cur_spread <= avg_spread - std_spread * context.trade_tolerance:
# if 0 or positive holdings in A, long spread -- long A, short B
return "long"
try:
log.info("No cash...")
except NameError:
log.info("avg_spread not defined...")
return None
def calculate_close_trades(context,signal,cur_spread, std_spread, a_order, b_order):
""" calculates close trades given close tolerance, signal, current spread, standard deviation of spread and current orders """
close_order_amount = np.matrix((-1*a_order, -1*b_order))
if signal =="long":
close_spread_amount = cur_spread + std_spread*context.close_tolerance
elif signal == "short":
close_spread_amount = cur_spread - std_spread*context.close_tolerance
if cur_spread in context.close_trades[signal]:
context.close_trades[signal][close_spread_amount] += close_order_amount
else:
context.close_trades[signal][close_spread_amount] = close_order_amount
return None
def calculate_net_orders(context, cur_spread, a_order, b_order):
""" calculates net orders given the close orders from existing trades, current spread and current orders """
for close_spread in context.close_trades["long"].keys():
if cur_spread > close_spread:
amount = context.close_trades["long"][close_spread]
a_order += amount.item(0)
b_order += amount.item(1)
del context.close_trades["long"][close_spread]
for close_spread in context.close_trades["short"].keys():
if cur_spread < close_spread:
amount = context.close_trades["short"][close_spread]
a_order += amount.item(0)
b_order += amount.item(1)
del context.close_trades["short"][close_spread]
return (a_order, b_order)
def calculate_order(context, signal, a_price, b_price):
""" returns the net-flat order amounts of KO and PEP given signal and price """
a_abs_order_amount = int(context.bet_amount/a_price)
b_abs_order_amount = int(context.bet_amount/b_price)
if not signal: return (0,0)
elif signal[0] == "l":
return ( a_abs_order_amount, (-1 * b_abs_order_amount))
elif signal[0] == "s":
return ( (-1 * a_abs_order_amount), b_abs_order_amount)
return None
@batch_transform(refresh_period=1, window_length=20)
def avgstd(datapanel,context):
""" returns mean and standard devaition of KO/PEP over the last 20 days"""
a_price = np.array(datapanel['price'][context.stocks[0]])
b_price = np.array(datapanel['price'][context.stocks[1]])
if a_price is not None and b_price is not None:
spread = a_price/b_price
return (spread.mean(), spread.std())
else:
return None