File size: 6,747 Bytes
6c3f19f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | import torch
import gpytorch
import math
from gpytorch.models import ExactGP
from gpytorch import settings as gptsettings
from gpytorch.constraints import GreaterThan,Positive
from gpytorch.distributions import MultivariateNormal
import gpytorch.kernels as kernels
from .horseshoe import LogHalfHorseshoePrior
from .mollified_uniform import MollifiedUniformPrior
from gpytorch.priors import NormalPrior,LogNormalPrior
from onescience.utils.GP_TO.transforms import softplus,inv_softplus
from typing import List,Tuple,Union
from typing import List, Union
class GPR(ExactGP):
"""Standard GP regression module for numerical inputs
:param train_x: The training inputs (size N x d). All input variables are expected
to be numerical. For best performance, scale the variables to the unit hypercube.
:type train_x: torch.Tensor
:param train_y: The training targets (size N)
:type train_y: torch.Tensor
:param correlation_kernel: Either a `gpytorch.kernels.Kernel` instance or one of the
following strings - 'RBFKernel' (radial basis kernel), 'Matern52Kernel' (twice
differentiable Matern kernel), 'Matern32Kernel' (first order differentiable Matern
kernel). If the former is specified, any hyperparameters to be estimated need to have
associated priors for multi-start optimization. If the latter is specified, then
the kernel uses a separate lengthscale for each input variable.
:type correlation_kernel: Union[gpytorch.kernels.Kernel,str]
:param noise: The (initial) noise variance.
:type noise: float, optional
:param fix_noise: Fixes the noise variance at the current level if `True` is specifed.
Defaults to `False`
:type fix_noise: bool, optional
:param lb_noise: Lower bound on the noise variance. Setting a higher value results in
more stable computations, when optimizing noise variance, but might reduce
prediction quality. Defaults to 1e-6
:type lb_noise: float, optional
"""
def __init__(
self,
train_x:torch.Tensor,
train_y:torch.Tensor,
correlation_kernel,
noise_indices:List[int],
noise:float=1e-4,
fix_noise:bool=False,
lb_noise:float=1e-12,
) -> None:
# check inputs
if not torch.is_tensor(train_x):
raise RuntimeError("'train_x' must be a tensor")
if not torch.is_tensor(train_y):
raise RuntimeError("'train_y' must be a tensor")
if train_x.shape[0] != train_y.shape[0]:
raise RuntimeError("Inputs and output have different number of observations")
# initializing likelihood
noise_constraint=GreaterThan(lb_noise,transform=torch.exp,inv_transform=torch.log)
if len(noise_indices) == 0:
likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=noise_constraint)
y_mean = torch.tensor(0.0)
y_std = torch.tensor(1.0)
train_y_sc = (train_y-y_mean)/y_std
ExactGP.__init__(self, train_x,train_y_sc, likelihood)
# registering mean and std of the raw response
self.register_buffer('y_mean',y_mean)
self.register_buffer('y_std',y_std)
self.register_buffer('y_scaled',train_y_sc)
self._num_outputs = 1
# initializing and fixing noise
if noise is not None:
self.likelihood.initialize(noise=noise)
self.likelihood.register_prior('noise_prior',LogHalfHorseshoePrior(0.01,lb_noise),'raw_noise')
if fix_noise:
self.likelihood.raw_noise.requires_grad_(False)
if isinstance(correlation_kernel,str):
try:
correlation_kernel_class = getattr(kernels,correlation_kernel)
correlation_kernel = correlation_kernel_class(
ard_num_dims = self.train_inputs[0].size(1),
lengthscale_constraint=Positive(transform=torch.exp,inv_transform=torch.log),
)
correlation_kernel.register_prior(
'lengthscale_prior',MollifiedUniformPrior(math.log(0.1),math.log(10)),'raw_lengthscale'
)
except:
raise RuntimeError(
"%s not an allowed kernel" % correlation_kernel
)
elif not isinstance(correlation_kernel,gpytorch.kernels.Kernel):
raise RuntimeError(
"specified correlation kernel is not a `gpytorch.kernels.Kernel` instance"
)
self.covar_module = kernels.ScaleKernel(
base_kernel = correlation_kernel,
outputscale_constraint=Positive(transform=softplus,inv_transform=inv_softplus),
)
# register priors
self.covar_module.register_prior(
'outputscale_prior',LogNormalPrior(1e-6,1.),'outputscale'
)
def forward(self,x:torch.Tensor)->MultivariateNormal:
mean_x = self.mean_module(x)
covar_x = self.covar_module(x)
return MultivariateNormal(mean_x,covar_x)
def predict(
self,x:torch.Tensor,return_std:bool=False,include_noise:bool=False
)-> Union[torch.Tensor,Tuple[torch.Tensor]]:
"""Returns the predictive mean, and optionally the standard deviation at the given points
:param x: The input variables at which the predictions are sought.
:type x: torch.Tensor
:param return_std: Standard deviation is returned along the predictions if `True`.
Defaults to `False`.
:type return_std: bool, optional
:param include_noise: Noise variance is included in the standard deviation if `True`.
Defaults to `False`.
:type include_noise: bool
"""
self.eval()
with gptsettings.fast_computations(log_prob=False):
# determine if batched or not
ndim = self.train_targets.ndim
if ndim == 1:
output = self(x)
else:
# for batched GPs
num_samples = self.train_targets.shape[0]
output = self(x.unsqueeze(0).repeat(num_samples,1,1))
if return_std and include_noise:
output = self.likelihood(output)
out_mean = self.y_mean + self.y_std*output.mean
# standard deviation may not always be needed
if return_std:
out_std = output.variance.sqrt()*self.y_std
return out_mean,out_std
return out_mean
|