mirror of
https://github.com/wassname/Volt.git
synced 2026-09-09 11:16:09 +08:00
298 lines
12 KiB
Python
298 lines
12 KiB
Python
import matplotlib.pyplot as plt
|
|
import seaborn as sns
|
|
import numpy as np
|
|
import torch
|
|
import pandas as pd
|
|
import os
|
|
import gpytorch
|
|
import argparse
|
|
import datetime
|
|
|
|
from botorch.models import SingleTaskGP
|
|
from botorch.optim.fit import fit_gpytorch_torch
|
|
from gpytorch.likelihoods import GaussianLikelihood
|
|
from gpytorch.mlls import ExactMarginalLogLikelihood
|
|
from gpytorch.means import ConstantMean, LinearMean
|
|
from gpytorch.kernels import SpectralMixtureKernel, MaternKernel, RBFKernel, ScaleKernel
|
|
from voltron.means import EWMAMean, DEWMAMean, TEWMAMean
|
|
from voltron.train_utils import LearnGPCV, TrainVolModel, TrainVoltMagpieModel, TrainBasicModel
|
|
from voltron.models import VoltMagpie
|
|
from voltron.means import LogLinearMean
|
|
|
|
from voltron.rollout_utils import GeneratePrediction, Rollouts
|
|
from voltron.data import make_ticker_list, DataGetter, GetStockHistory
|
|
|
|
|
|
def GenerateGPCVPredictions(ticker, dat,
|
|
forecast_horizon=20, ntimes=25,
|
|
train_iters=400, nsample=1000,
|
|
ntrain=400):
|
|
|
|
end_idxs = torch.arange(ntrain, dat.shape[0],
|
|
int((dat.shape[0]-ntrain)/ntimes))
|
|
ntest = forecast_horizon
|
|
dt = 1./252
|
|
|
|
savepath = "./saved-outputs/" + ticker + "/"
|
|
if not os.path.exists(savepath):
|
|
os.mkdir(savepath)
|
|
for last_day in end_idxs:
|
|
date = str(dat.index[last_day.item()].date())
|
|
print(date, ticker)
|
|
train_y = torch.FloatTensor(dat.Close[last_day.item()-ntrain:last_day.item()].to_numpy())
|
|
train_x = torch.arange(train_y.shape[0]-1) * dt
|
|
test_x = torch.arange(ntest) * dt + train_x[-1] + train_x[1]
|
|
# try:
|
|
use_cuda = torch.cuda.is_available()
|
|
if use_cuda:
|
|
train_x = train_x.cuda()
|
|
test_x = test_x.cuda()
|
|
train_y = train_y.cuda()
|
|
|
|
model, likelihood = LearnGPCV(train_x, train_y,
|
|
train_iters=train_iters, printing=False, return_model=True)
|
|
preds = likelihood(model(test_x),
|
|
return_gaussian=False).sample(torch.Size((nsample,)))
|
|
preds = preds.cumsum(-1).squeeze()
|
|
preds = preds.view(-1, preds.shape[-1])
|
|
save_samples = preds.view(-1, preds.shape[-1]) * (dt ** 0.5) + train_y[-1].log()
|
|
torch.save(save_samples, savepath + "gpcv_" + date + ".pt")
|
|
|
|
return
|
|
|
|
def GenerateStockPredictions(ticker, dat,
|
|
forecast_horizon=20,
|
|
train_iters=400, nsample=1000,
|
|
ntrain=400, mean='ewma', kernel='volt',
|
|
save=False, k=300, ntimes=-1):
|
|
|
|
if ntimes == -1:
|
|
end_idxs = torch.arange(ntrain, dat.shape[0])
|
|
else:
|
|
end_idxs = torch.arange(ntrain, dat.shape[0],
|
|
int((dat.shape[0]-ntrain)/ntimes))
|
|
ntest = forecast_horizon
|
|
dt = 1./252
|
|
|
|
model_name = kernel + "_" + mean + str(k) + "_"
|
|
par_dir = "./saved-outputs/"
|
|
if not os.path.exists(par_dir):
|
|
os.mkdir(par_dir)
|
|
savepath = par_dir + ticker + "/"
|
|
if not os.path.exists(savepath):
|
|
os.mkdir(savepath)
|
|
|
|
for last_day in end_idxs:
|
|
date = str(dat.index[last_day.item()].date())
|
|
# try:
|
|
train_y = torch.FloatTensor(dat.Close[last_day.item()-ntrain:last_day.item()].to_numpy())
|
|
train_x = torch.arange(train_y.shape[0]-1) * dt
|
|
test_x = torch.arange(ntest) * dt + train_x[-1] + train_x[1]
|
|
# try:
|
|
use_cuda = torch.cuda.is_available()
|
|
if use_cuda:
|
|
train_x = train_x.cuda()
|
|
test_x = test_x.cuda()
|
|
train_y = train_y.cuda()
|
|
|
|
# print("Producing " + ticker + " Forecasts.....")
|
|
if kernel == "volt":
|
|
vol = LearnGPCV(train_x, train_y, train_iters=train_iters,
|
|
printing=False)
|
|
vmod, vlh = TrainVolModel(train_x, vol,
|
|
train_iters=train_iters, printing=False)
|
|
voltron, lh = TrainVoltMagpieModel(train_x, train_y[1:],
|
|
vmod, vlh, vol,
|
|
printing=False,
|
|
train_iters=train_iters,
|
|
k=k, mean_func=mean)
|
|
vmod.eval();
|
|
if mean in ['ewma', 'dewma', 'tewma']:
|
|
save_samples = Rollouts(train_x, train_y, test_x, voltron,
|
|
nsample=nsample)
|
|
|
|
else: ## VOLT + STANDARD MEAN
|
|
voltron.vol_model.eval()
|
|
predvol = voltron.vol_model(test_x).sample(torch.Size((nsample, ))).exp()
|
|
save_samples = GeneratePrediction(train_x, train_y, test_x,
|
|
predvol, voltron).detach()
|
|
del predvol
|
|
|
|
del voltron, lh, vmod, vlh, vol
|
|
torch.cuda.empty_cache()
|
|
|
|
if save:
|
|
if not os.path.exists(savepath):
|
|
os.mkdir(savepath)
|
|
|
|
torch.save(save_samples, savepath + model_name + date + ".pt")
|
|
# except:
|
|
# nans = torch.ones(nsample, ntest) * torch.nan
|
|
# if save:
|
|
# if not os.path.exists(savepath):
|
|
# os.mkdir(savepath)
|
|
# torch.save(nans, savepath + model_name + date + ".pt")
|
|
|
|
|
|
return dat, save_samples
|
|
|
|
|
|
|
|
def GenerateOneDayPredictions(ticker, train_y, date,
|
|
forecast_horizon=20,
|
|
train_iters=400, nsample=1000,
|
|
ntrain=400, save=False, mean=None):
|
|
|
|
ntest = forecast_horizon
|
|
dt = 1./252
|
|
par_dir = "./saved-outputs/"
|
|
if not os.path.exists(par_dir):
|
|
os.mkdir(par_dir)
|
|
savepath = par_dir + ticker + "/"
|
|
if not os.path.exists(savepath):
|
|
os.mkdir(savepath)
|
|
|
|
train_x = torch.arange(train_y.shape[0]-1) * dt
|
|
test_x = torch.arange(ntest) * dt + train_x[-1] + train_x[1]
|
|
use_cuda = torch.cuda.is_available()
|
|
if use_cuda:
|
|
train_x = train_x.cuda()
|
|
test_x = test_x.cuda()
|
|
train_y = train_y.cuda()
|
|
|
|
vol = LearnGPCV(train_x, train_y, train_iters=train_iters,
|
|
printing=False)
|
|
vmod, vlh = TrainVolModel(train_x, vol,
|
|
train_iters=train_iters, printing=False)
|
|
|
|
if mean=='constant':
|
|
voltron, lh = TrainVoltMagpieModel(train_x, train_y[1:],
|
|
vmod, vlh, vol,
|
|
printing=False,
|
|
train_iters=200,
|
|
mean_func='constant')
|
|
vmod.eval();
|
|
voltron.eval();
|
|
save_samples = Rollouts(train_x, train_y, test_x, voltron,
|
|
nsample=nsample)
|
|
|
|
if save:
|
|
model_name = "volt_" + mean + "_"
|
|
torch.save(save_samples, savepath + model_name + date + ".pt")
|
|
else:
|
|
for mean in ['ewma', 'dewma', 'tewma']:
|
|
for k in [25, 50, 100, 200, 300, 400]:
|
|
try:
|
|
voltron, lh = TrainVoltMagpieModel(train_x, train_y[1:],
|
|
vmod, vlh, vol,
|
|
printing=False,
|
|
train_iters=0,
|
|
k=k, mean_func=mean)
|
|
vmod.eval();
|
|
voltron.eval();
|
|
save_samples = Rollouts(train_x, train_y, test_x, voltron,
|
|
nsample=nsample)
|
|
except:
|
|
print("Failed: ", ticker, mean, k)
|
|
if save:
|
|
save_samples = torch.ones(nsample, ntest) * torch.nan
|
|
|
|
if save:
|
|
model_name = "volt_" + mean + str(k) + "_"
|
|
torch.save(save_samples, savepath + model_name + date + ".pt")
|
|
|
|
del voltron, lh, vmod, vlh, vol, save_samples
|
|
torch.cuda.empty_cache()
|
|
return
|
|
|
|
|
|
|
|
def GenerateBasicPredictions(ticker, dat, kernel_name, mean_name='ewma', k=400,
|
|
forecast_horizon=100,
|
|
train_iters=600, nsample=1000,
|
|
ntrain=400, save=False, ntimes=-1):
|
|
|
|
if ntimes == -1:
|
|
end_idxs = torch.arange(ntrain, dat.shape[0])
|
|
else:
|
|
end_idxs = torch.arange(ntrain, dat.shape[0],
|
|
int((dat.shape[0]-ntrain)/ntimes))
|
|
ntest = forecast_horizon
|
|
dt = 1./252
|
|
par_dir = "./saved-outputs/"
|
|
if not os.path.exists(par_dir):
|
|
os.mkdir(par_dir)
|
|
savepath = par_dir + ticker + "/"
|
|
if not os.path.exists(savepath):
|
|
os.mkdir(savepath)
|
|
|
|
for last_day in end_idxs:
|
|
date = str(dat.index[last_day.item()].date())
|
|
train_y = torch.FloatTensor(dat.Close[last_day.item()-ntrain:last_day.item()].to_numpy())
|
|
train_x = torch.arange(train_y.shape[0]-1) * dt
|
|
test_x = torch.arange(ntest) * dt + train_x[-1] + train_x[1]
|
|
# try:
|
|
use_cuda = torch.cuda.is_available()
|
|
if use_cuda:
|
|
train_x = train_x.cuda()
|
|
test_x = test_x.cuda()
|
|
train_y = train_y.cuda()
|
|
|
|
# print("Producing " + ticker + " Forecasts.....")
|
|
kernel_possibilities = {"sm": SpectralMixtureKernel,
|
|
"matern": MaternKernel,
|
|
"rbf": RBFKernel}
|
|
kernel = kernel_possibilities[kernel_name.lower()]
|
|
if kernel_name.lower() != 'sm':
|
|
kernel = ScaleKernel(kernel())
|
|
else:
|
|
kernel = kernel(num_mixtures=15)
|
|
kernel.initialize_from_data_empspect(train_x, train_y.log())
|
|
|
|
train_y = train_y[1:]
|
|
|
|
model = SingleTaskGP(
|
|
train_x.view(-1,1),
|
|
train_y.log().reshape(-1, 1),
|
|
covar_module=kernel,
|
|
likelihood=GaussianLikelihood()
|
|
)
|
|
|
|
mean_name = mean_name.lower()
|
|
if mean_name == "loglinear":
|
|
model.mean_module = LogLinearMean(1)
|
|
model.mean_module.initialize_from_data(train_x, train_y.log())
|
|
elif mean_name == 'linear':
|
|
model.mean_module = LinearMean(1)
|
|
elif mean_name == "constant":
|
|
model.mean_module = ConstantMean()
|
|
elif mean_name == "ewma":
|
|
model.mean_module = EWMAMean(train_x, train_y.log(), k=k).to(train_x.device)
|
|
elif mean_name == "dewma":
|
|
model.mean_module = DEWMAMean(train_x, train_y.log(), k=k).to(train_x.device)
|
|
elif mean_name == "tewma":
|
|
model.mean_module = TEWMAMean(train_x, train_y.log(), k=k).to(train_x.device)
|
|
|
|
if use_cuda:
|
|
model = model.to(train_x.device)
|
|
print("Fitting Model", ticker)
|
|
mll = ExactMarginalLogLikelihood(model.likelihood, model)
|
|
fit_gpytorch_torch(mll, options={'maxiter':train_iters, 'disp':False})
|
|
|
|
if mean_name in ["loglinear", "constant", 'linear']:
|
|
save_samples = model.posterior(test_x).sample(torch.Size((nsample,
|
|
))).squeeze(-1).cpu().detach()
|
|
else:
|
|
save_samples = Rollouts(
|
|
train_x, train_y, test_x, model, nsample=nsample, method = "nonvol"
|
|
).cpu().detach()
|
|
|
|
model_name = kernel_name + "_" + mean_name + str(k) + "_"
|
|
torch.save(save_samples, savepath + model_name + date + ".pt")
|
|
|
|
model.train()
|
|
torch.cuda.empty_cache()
|
|
del model
|
|
|
|
|
|
return dat, save_samples |