Files
2022-07-17 20:44:16 +08:00

258 lines
9.6 KiB
Python

import numpy as np
import torch
import gpytorch
import sys
sys.path.append("../")
from voltron.likelihoods import VolatilityGaussianLikelihood
from voltron.models import SingleTaskVariationalGP
from voltron.kernels import BMKernel, VolatilityKernel, FBMKernel
from voltron.models import BMGP, VoltronGP, MaternGP, SMGP, VoltMagpie
from voltron.means import LogLinearMean, EWMAMean, DEWMAMean, TEWMAMean, MeanRevertingEMAMean
from gpytorch.kernels import ScaleKernel, RBFKernel, MaternKernel
def LearnGPCV(train_x, train_y, train_iters=1000, printing=False, early_stopping=False, kernel = "bm"):
dt = train_x[1]-train_x[0]
scaled_returns = (train_y[1:] - train_y[:-1]) / (train_y[:-1]) / (dt**0.5)
yy = scaled_returns
likelihood = VolatilityGaussianLikelihood(param="exp")
# likelihood.raw_a.data -= 4.
if kernel == "bm":
covar_module = BMKernel()
elif kernel == "fbm":
covar_module = FBMKernel()
model = SingleTaskVariationalGP(
init_points=train_x.view(-1,1), likelihood=likelihood, use_piv_chol_init=False,
mean_module = gpytorch.means.ConstantMean(), covar_module=covar_module,
learn_inducing_locations=False, use_whitened_var_strat=False
)
model.initialize_variational_parameters(likelihood, train_x, y=yy)
model.train()
likelihood.train()
# Use the adam optimizer
optimizer = torch.optim.Adam([
{"params": model.parameters()},
# {"params": likelihood.parameters(), "lr": 0.1}
], lr=0.01)
# "Loss" for GPs - the marginal log likelihood
# num_data refers to the number of training datapoints
mll = gpytorch.mlls.VariationalELBO(likelihood, model, yy.numel(), combine_terms = True)
old_loss = 10000.
print_every = 50
for i in range(train_iters):
# Zero backpropped gradients from previous iteration
optimizer.zero_grad()
# Get predictive output
with gpytorch.settings.num_gauss_hermite_locs(75):
output = model(train_x)
# Calc loss and backprop gradients
loss = -mll(output, yy)
loss.backward()
if printing:
if i % print_every == 0:
print('Iter %d/%d - Loss: %.3f' % (i + 1, train_iters, loss.item()))
optimizer.step()
model.eval();
likelihood.eval();
predictive = model(train_x)
pred_scale = likelihood(predictive, return_gaussian=False).scale.mean(0).detach()
return pred_scale
def TrainVolModel(train_x, vol_path, train_iters=1000, printing=False, kernel = "bm", **kwargs):
vol_lh = gpytorch.likelihoods.GaussianLikelihood().to(train_x.device)
vol_lh.noise.data = torch.tensor([1e-2])
vol_model = BMGP(train_x, vol_path.log(), vol_lh, kernel=kernel, **kwargs).to(train_x.device)
# vol_model.covar_module.raw_vol.data = torch.tensor([-3.])
optimizer = torch.optim.Adam([
{'params': vol_model.parameters()}, # Includes GaussianLikelihood parameters
], lr=0.01)
# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(vol_lh, vol_model)
print_every = 50
for i in range(train_iters):
# Zero gradients from previous iteration
optimizer.zero_grad()
# Output from model
output = vol_model(train_x)
# Calc loss and backprop gradients
loss = -mll(output, vol_path.log())
loss.backward()
if printing:
if i % print_every == 0:
print('Iter %d/%d - Loss: %.3f' % (i + 1, train_iters, loss.item()))
optimizer.step()
return vol_model, vol_lh
def TrainDataModel(train_x, train_y, vol_model, vol_lh, vol_path,
train_iters=1000, printing=False):
voltron_lh = gpytorch.likelihoods.GaussianLikelihood()
voltron = VoltronGP(train_x, train_y.log(), voltron_lh, vol_path)
voltron.mean_module = LogLinearMean(1)
# voltron.mean_module.register_prior("slope_prior",
# gpytorch.priors.NormalPrior(0,0.1), 'weights')
voltron.mean_module.initialize_from_data(train_x, train_y.log())
# voltron.mean_module.weights.data = torch.tensor([[0.]])
voltron.likelihood.raw_noise.data = torch.tensor([1e-5])
voltron.vol_lh = vol_lh
voltron.vol_model = vol_model
grad_flags = [True, True, True, False, False, False]
for idx, p in enumerate(voltron.parameters()):
p.requires_grad = grad_flags[idx]
voltron.train();
voltron_lh.train();
voltron.vol_lh.train();
voltron.vol_model.train();
# Use the adam optimizer
optimizer = torch.optim.Adam([
{'params': voltron.parameters()}, # Includes GaussianLikelihood parameters
], lr=0.1)
# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(voltron_lh, voltron)
print_every = 50
for i in range(train_iters):
# Zero gradients from previous iteration
optimizer.zero_grad()
# Output from model
output = voltron(train_x)
# Calc loss and backprop gradients
loss = -mll(output, train_y.log())
loss.backward()
if printing:
if i % print_every == 0:
print('Iter %d/%d - Loss: %.3f' % (i + 1, train_iters, loss.item()))
optimizer.step()
return voltron, voltron_lh
def TrainBasicModel(train_x, train_y, train_iters=1000, printing=False,
model_type="matern",
num_mixtures=10, mean_func="loglinear"):
lh = gpytorch.likelihoods.GaussianLikelihood()
if model_type == "matern":
model = MaternGP(train_x, train_y.log(), lh)
else:
model = SMGP(train_x, train_y.log(), lh, num_mixtures)
if mean_func == "loglinear":
model.mean_module = LogLinearMean(1)
model.mean_module.register_prior("slope_prior",
gpytorch.priors.NormalPrior(0,0.1), 'weights')
model.mean_module.initialize_from_data(train_x, train_y.log())
model.likelihood.raw_noise.data = torch.tensor([1e-5])
model = model.to(train_x.device)
model.train();
lh.train();
# Use the adam optimizer
optimizer = torch.optim.Adam([
{'params': model.parameters()}, # Includes GaussianLikelihood parameters
], lr=0.1)
# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(lh, model)
print_every = 50
for i in range(train_iters):
# Zero gradients from previous iteration
optimizer.zero_grad()
# Output from model
output = model(train_x)
# Calc loss and backprop gradients
loss = -mll(output, train_y.log())
loss.backward()
if printing:
if i % print_every == 0:
print('Iter %d/%d - Loss: %.3f' % (i + 1, train_iters, loss.item()))
optimizer.step()
return model, lh
def TrainVoltMagpieModel(train_x, train_y, vol_model, vol_lh, vol_path,
train_iters=1000, printing=False, k=25, theta=0.5,
mean_func="ewma"):
voltron_lh = gpytorch.likelihoods.GaussianLikelihood().to(train_x.device)
voltron = VoltMagpie(train_x, train_y.log(),
voltron_lh, vol_path, k=k).to(train_x.device)
if mean_func.lower() in ["ewma", "dewma", "tewma", "meanrevert"]:
# default voltmagpie is an ewma mean so we don't need to redefine anything
grad_flags = [True, False, False, False]
if mean_func.lower() == "dewma":
voltron.mean_module = DEWMAMean(train_x, train_y.log(), k).to(train_x.device)
elif mean_func.lower() == 'tewma':
voltron.mean_module = TEWMAMean(train_x, train_y.log(), k).to(train_x.device)
elif mean_func.lower() == 'meanrevert':
voltron.mean_module = MeanRevertingEMAMean(train_x, train_y.log(),
k, theta).to(train_x.device)
elif mean_func.lower()=='constant':
voltron.mean_module = gpytorch.means.ConstantMean().to(train_x.device)
grad_flags = [True, True, False, False, False]
elif mean_func.lower()=='loglinear':
voltron.mean_module = LogLinearMean(1).to(train_x.device)
voltron.mean_module.initialize_from_data(train_x, train_y.log())
grad_flags = [True, True, True, False, False, False]
elif mean_func.lower()=='linear':
voltron.mean_module = gpytorch.means.LinearMean(1).to(train_x.device)
grad_flags = [True, True, True, False, False, False]
voltron.likelihood.raw_noise.data = torch.tensor([1e-5]).to(train_x.device)
voltron.vol_lh = vol_lh.to(train_x.device)
voltron.vol_model = vol_model.to(train_x.device)
for idx, p in enumerate(voltron.parameters()):
p.requires_grad = grad_flags[idx]
voltron.train();
voltron_lh.train();
voltron.vol_lh.train();
voltron.vol_model.train();
# Use the adam optimizer
optimizer = torch.optim.Adam([
{'params': voltron.parameters()}, # Includes GaussianLikelihood parameters
], lr=0.1)
# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(voltron_lh, voltron)
print_every = 50
for i in range(train_iters):
# Zero gradients from previous iteration
optimizer.zero_grad()
# Output from model
output = voltron(train_x)
# Calc loss and backprop gradients
loss = -mll(output, train_y.log())
loss.backward()
if printing:
if i % print_every == 0:
print('Iter %d/%d - Loss: %.3f' % (i + 1, train_iters, loss.item()))
optimizer.step()
return voltron, voltron_lh