95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
import datetime as dt
|
|
import pandas as pd
|
|
import util
|
|
import indicators
|
|
from BagLearner import BagLearner
|
|
from RTLearner import RTLearner
|
|
|
|
|
|
class StrategyLearner(object):
|
|
|
|
def __init__(self, verbose=False, impact=0.0, commission=0.0, testing=False):
|
|
self.verbose = verbose
|
|
self.impact = impact
|
|
self.commission = commission
|
|
self.testing = testing
|
|
|
|
def _get_volume(self):
|
|
"""For reference."""
|
|
volume_all = ut.get_data(syms, dates, colname="Volume")
|
|
volume = volume_all[syms] # only portfolio symbols
|
|
# volume_SPY = volume_all['SPY'] # only SPY, for comparison later
|
|
if self.verbose:
|
|
print(volume)
|
|
|
|
def _add_indicators(self, df, symbol):
|
|
"""Add indicators for learning to DataFrame."""
|
|
df.drop(columns=["SPY"], inplace=True)
|
|
indicators.macd(df, symbol)
|
|
indicators.rsi(df, symbol)
|
|
indicators.price_sma(df, symbol, [8])
|
|
indicators.price_delta(df, symbol, 3)
|
|
df.dropna(inplace=True)
|
|
|
|
def addEvidence(self, symbol="IBM",
|
|
sd=dt.datetime(2008, 1, 1),
|
|
ed=dt.datetime(2009, 1, 1),
|
|
sv=10000):
|
|
|
|
self.indicators = ['macd_diff', 'rsi', 'price_sma_8']
|
|
df = util.get_data([symbol], pd.date_range(sd, ed))
|
|
self._add_indicators(df, symbol)
|
|
|
|
def classify_y(row):
|
|
if row > 0.1:
|
|
return 1
|
|
elif row < -0.1:
|
|
return -1
|
|
return 0
|
|
|
|
self.learner = RTLearner(leaf_size = 7)
|
|
# self.learner = BagLearner(RTLearner, 5, {'leaf_size': 5})
|
|
data_x = df[self.indicators].to_numpy()
|
|
y = df['pct_3'].apply(classify_y)
|
|
self.learner.addEvidence(data_x, y.to_numpy())
|
|
return y
|
|
|
|
def strat(self, data_y, orders):
|
|
self.holding = 0
|
|
|
|
def strat(row):
|
|
y = int(data_y.loc[row.name][0])
|
|
shares = 0
|
|
if self.holding == 0 and y == 1:
|
|
shares = 1000
|
|
elif self.holding == -1000 and y == 1:
|
|
shares = 2000
|
|
elif self.holding == 0 and y == -1:
|
|
shares = -1000
|
|
elif self.holding == 1000 and y == -1:
|
|
shares = -2000
|
|
self.holding += shares
|
|
return shares
|
|
|
|
orders["Shares"] = orders.apply(strat, axis=1)
|
|
|
|
def testPolicy(self, symbol="IBM",
|
|
sd=dt.datetime(2009, 1, 1),
|
|
ed=dt.datetime(2010, 1, 1),
|
|
sv=10000):
|
|
df = util.get_data([symbol], pd.date_range(sd, ed))
|
|
self._add_indicators(df, symbol)
|
|
data_x = df[self.indicators].to_numpy()
|
|
data_y = pd.DataFrame(index=df.index, data=self.learner.query(data_x))
|
|
|
|
orders = pd.DataFrame(index=df.index)
|
|
orders["Symbol"] = symbol
|
|
orders["Order"] = ""
|
|
orders["Shares"] = 0
|
|
self.strat(data_y, orders)
|
|
if self.testing:
|
|
return orders
|
|
else:
|
|
return orders[["Shares"]]
|
|
|