| 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:
|
|
|
| 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")
|
|
|
|
|
| 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)
|
|
|
|
|
| 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
|
|
|
|
|
| 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),
|
| )
|
|
|
| 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):
|
|
|
| ndim = self.train_targets.ndim
|
| if ndim == 1:
|
| output = self(x)
|
| else:
|
|
|
| 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
|
|
|
|
|
| if return_std:
|
| out_std = output.variance.sqrt()*self.y_std
|
| return out_mean,out_std
|
|
|
| return out_mean
|
|
|
