65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
import pandas as pd
|
|
from util import get_data
|
|
from collections import namedtuple
|
|
|
|
Position = namedtuple("Pos", ["cash", "shares", "transactions"])
|
|
|
|
|
|
def author():
|
|
return "felixm"
|
|
|
|
|
|
def new_positions(positions, price):
|
|
"""Calculate all potential new positions and then keep the best three."""
|
|
# Execute all possible transactions
|
|
new_positions = []
|
|
for p in positions:
|
|
for t in [-2000, -1000, 0, 1000, 2000]:
|
|
ts = p.transactions + [t]
|
|
p_new = Position(p.cash - t * price, p.shares + t, ts)
|
|
new_positions.append(p_new)
|
|
|
|
# Keep the positions with the highest cash value for each amount of shares.
|
|
best = {}
|
|
for p in new_positions:
|
|
if p.shares not in [-1000, 0, 1000]:
|
|
pass
|
|
elif p.shares in best and p.cash > best[p.shares].cash:
|
|
best[p.shares] = p
|
|
elif not p.shares in best:
|
|
best[p.shares] = p
|
|
return list(best.values())
|
|
|
|
|
|
def transactions_to_orders(transactions, prices, symbol):
|
|
order = pd.Series("", index=prices.index)
|
|
shares = pd.Series(0, index=prices.index)
|
|
|
|
for i, t in enumerate(transactions):
|
|
if t > 0:
|
|
order.iloc[i] = "BUY"
|
|
shares.iloc[i] = t
|
|
if t < 0:
|
|
order.iloc[i] = "SELL"
|
|
shares.iloc[i] = -t
|
|
|
|
prices["Symbol"] = pd.Series(symbol, index=prices.index)
|
|
prices["Order"] = order
|
|
prices["Shares"] = shares
|
|
prices.drop(columns=[symbol], inplace=True)
|
|
prices = prices[shares != 0]
|
|
return prices
|
|
|
|
|
|
def testPolicy(symbol, sd, ed, sv):
|
|
prices = get_data([symbol], pd.date_range(sd, ed))
|
|
prices.drop(columns=["SPY"], inplace=True)
|
|
|
|
positions = [Position(sv, 0, [])]
|
|
for date, price in prices.iterrows():
|
|
positions = new_positions(positions, price[0])
|
|
|
|
price = prices.iloc[-1][symbol]
|
|
best_position = max(positions, key=lambda p: p.cash + p.shares * price)
|
|
return transactions_to_orders(best_position.transactions, prices, symbol)
|