def create_kernel_no_ard(self, **kwargs):
return MaternKernel(nu=2.5, **kwargs)
def create_kernel_ard(self, num_dims, **kwargs):
return MaternKernel(nu=1.5, ard_num_dims=num_dims, **kwargs)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
likelihood: Optional[Likelihood] = None,
covar_modules: Optional[List[Kernel]] = None,
num_latent_dims: Optional[List[int]] = None,
learn_latent_pars: bool = True,
latent_init: str = "default",
outcome_transform: Optional[OutcomeTransform] = None,
input_transform: Optional[InputTransform] = None,
):
r"""A HigherOrderGP model for high-dim output regression.
Args:
train_X: A `batch_shape x n x d`-dim tensor of training inputs.
train_Y: A `batch_shape x n x output_shape`-dim tensor of training targets.
likelihood: Gaussian likelihood for the model.
covar_modules: List of kernels for each output structure.
num_latent_dims: Sizes for the latent dimensions.
learn_latent_pars: If true, learn the latent parameters.
latent_init: [default or gp] how to initialize the latent parameters.
"""
if input_transform is not None:
input_transform.to(train_X)
# infer the dimension of `output_shape`.
num_output_dims = train_Y.dim() - train_X.dim() + 1
batch_shape = train_X.shape[:-2]
if len(batch_shape) > 1:
raise NotImplementedError(
"HigherOrderGP currently only supports 1-dim `batch_shape`.")
if outcome_transform is not None:
if isinstance(outcome_transform, Standardize) and not isinstance(
outcome_transform, FlattenedStandardize):
warnings.warn(
"HigherOrderGP does not support the outcome_transform "
"`Standardize`! Using `FlattenedStandardize` with `output_shape="
f"{train_Y.shape[- num_output_dims:]} and batch_shape="
f"{batch_shape} instead.",
RuntimeWarning,
)
outcome_transform = FlattenedStandardize(
output_shape=train_Y.shape[-num_output_dims:],
batch_shape=batch_shape,
)
train_Y, _ = outcome_transform(train_Y)
self._aug_batch_shape = batch_shape
self._num_dimensions = num_output_dims + 1
self._num_outputs = train_Y.shape[0] if batch_shape else 1
self.target_shape = train_Y.shape[-num_output_dims:]
self._input_batch_shape = batch_shape
if likelihood is None:
noise_prior = GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration -
1) / noise_prior.rate
likelihood = GaussianLikelihood(
noise_prior=noise_prior,
batch_shape=self._aug_batch_shape,
noise_constraint=GreaterThan(
MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=noise_prior_mode,
),
)
else:
self._is_custom_likelihood = True
super().__init__(
train_X,
train_Y.view(*self._aug_batch_shape, -1),
likelihood=likelihood,
)
if covar_modules is not None:
self.covar_modules = ModuleList(covar_modules)
else:
self.covar_modules = ModuleList([
MaternKernel(
nu=2.5,
lengthscale_prior=GammaPrior(3.0, 6.0),
batch_shape=self._aug_batch_shape,
ard_num_dims=1 if dim > 0 else train_X.shape[-1],
) for dim in range(self._num_dimensions)
])
if num_latent_dims is None:
num_latent_dims = [1] * (self._num_dimensions - 1)
self.to(train_X)
self._initialize_latents(
latent_init=latent_init,
num_latent_dims=num_latent_dims,
learn_latent_pars=learn_latent_pars,
device=train_Y.device,
dtype=train_Y.dtype,
)
if outcome_transform is not None:
self.outcome_transform = outcome_transform
if input_transform is not None:
self.input_transform = input_transform
def test_random_fourier_features(self):
# test kernel that is not Scale, RBF, or Matern
with self.assertRaises(NotImplementedError):
RandomFourierFeatures(
kernel=PeriodicKernel(),
input_dim=2,
num_rff_features=3,
)
# test batched kernel
with self.assertRaises(NotImplementedError):
RandomFourierFeatures(
kernel=RBFKernel(batch_shape=torch.Size([2])),
input_dim=2,
num_rff_features=3,
)
tkwargs = {"device": self.device}
for dtype in (torch.float, torch.double):
tkwargs["dtype"] = dtype
# test init
# test ScaleKernel
base_kernel = RBFKernel(ard_num_dims=2)
kernel = ScaleKernel(base_kernel).to(**tkwargs)
rff = RandomFourierFeatures(
kernel=kernel,
input_dim=2,
num_rff_features=3,
)
self.assertTrue(torch.equal(rff.outputscale, kernel.outputscale))
# check that rff makes a copy
self.assertFalse(rff.outputscale is kernel.outputscale)
self.assertTrue(
torch.equal(rff.lengthscale, base_kernel.lengthscale))
# check that rff makes a copy
self.assertFalse(rff.lengthscale is kernel.lengthscale)
# test not ScaleKernel
rff = RandomFourierFeatures(
kernel=base_kernel,
input_dim=2,
num_rff_features=3,
)
self.assertTrue(
torch.equal(rff.outputscale, torch.tensor(1, **tkwargs)))
self.assertTrue(
torch.equal(rff.lengthscale, base_kernel.lengthscale))
# check that rff makes a copy
self.assertFalse(rff.lengthscale is kernel.lengthscale)
self.assertEqual(rff.weights.shape, torch.Size([2, 3]))
self.assertEqual(rff.bias.shape, torch.Size([3]))
self.assertTrue(((rff.bias <= 2 * pi) & (rff.bias >= 0.0)).all())
# test forward
rff = RandomFourierFeatures(
kernel=kernel,
input_dim=2,
num_rff_features=3,
)
for batch_shape in (torch.Size([]), torch.Size([3])):
X = torch.rand(*batch_shape, 1, 2, **tkwargs)
Y = rff(X)
self.assertTrue(Y.shape, torch.Size([*batch_shape, 1, 1]))
expected_Y = torch.sqrt(
2 * rff.outputscale / rff.weights.shape[-1]) * (torch.cos(
X / base_kernel.lengthscale @ rff.weights + rff.bias))
self.assertTrue(torch.equal(Y, expected_Y))
# test get_weights
with mock.patch("torch.randn", wraps=torch.randn) as mock_randn:
rff._get_weights(base_kernel=base_kernel,
input_dim=2,
num_rff_features=3)
mock_randn.assert_called_once_with(
2,
3,
dtype=base_kernel.lengthscale.dtype,
device=base_kernel.lengthscale.device,
)
# test get_weights with Matern kernel
with mock.patch("torch.randn",
wraps=torch.randn) as mock_randn, mock.patch(
"torch.distributions.Gamma",
wraps=torch.distributions.Gamma) as mock_gamma:
base_kernel = MaternKernel(ard_num_dims=2).to(**tkwargs)
rff._get_weights(base_kernel=base_kernel,
input_dim=2,
num_rff_features=3)
mock_randn.assert_called_once_with(
2,
3,
dtype=base_kernel.lengthscale.dtype,
device=base_kernel.lengthscale.device,
)
mock_gamma.assert_called_once_with(
base_kernel.nu,
base_kernel.nu,
)
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
likelihood: Optional[Likelihood] = None,
covar_modules: Optional[List[Kernel]] = None,
num_latent_dims: Optional[List[int]] = None,
learn_latent_pars: bool = True,
first_dim_is_batch: bool = False,
latent_init: str = "default",
outcome_transform: Optional[OutcomeTransform] = None,
input_transform: Optional[InputTransform] = None,
):
r"""A HigherOrderGP model for high-dim output regression.
Args:
train_X: Training inputs
train_Y: Training targets
likelihood: Gaussian likelihood for the model.
covar_modules: List of kernels for each output structure.
num_latent_dims: Sizes for the latent dimensions.
learn_latent_pars: If true, learn the latent parameters.
first_dim_is_batch: If true, the first dimension of train_Y should be
regarded as a batch dimension (e.g. predicting batches of tensors).
latent_init: [default or gp] how to initialize the latent parameters.
"""
if input_transform is not None:
input_transform.to(train_X)
if outcome_transform is not None:
train_Y, _ = outcome_transform(train_Y)
if first_dim_is_batch:
self._aug_batch_shape = train_Y.shape[:1]
self._num_dimensions = len(train_Y.shape) - 1
self._num_outputs = train_Y.shape[0]
else:
self._aug_batch_shape = Size()
self._num_dimensions = len(train_Y.shape)
self._num_outputs = 1
self._input_batch_shape = train_X.shape[:-2]
if likelihood is None:
noise_prior = GammaPrior(1.1, 0.05)
noise_prior_mode = (noise_prior.concentration - 1) / noise_prior.rate
likelihood = GaussianLikelihood(
noise_prior=noise_prior,
batch_shape=self._aug_batch_shape,
noise_constraint=GreaterThan(
MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=noise_prior_mode,
),
)
else:
self._is_custom_likelihood = True
super().__init__(
train_X,
train_Y.view(*self._aug_batch_shape, -1),
likelihood=likelihood,
)
if covar_modules is not None:
self.covar_modules = ModuleList(covar_modules)
else:
self.covar_modules = ModuleList(
[
MaternKernel(
nu=2.5,
lengthscale_prior=GammaPrior(3.0, 6.0),
batch_shape=self._aug_batch_shape,
ard_num_dims=1 if dim > 0 else train_X.shape[-1],
)
for dim in range(self._num_dimensions)
]
)
if num_latent_dims is None:
num_latent_dims = [1] * (self._num_dimensions - 1)
if first_dim_is_batch:
self.target_shape = train_Y.shape[2:]
else:
self.target_shape = train_Y.shape[1:]
self.to(train_X.device)
self._initialize_latents(
latent_init=latent_init,
num_latent_dims=num_latent_dims,
learn_latent_pars=learn_latent_pars,
device=train_Y.device,
dtype=train_Y.dtype,
)
if outcome_transform is not None:
self.outcome_transform = outcome_transform
if input_transform is not None:
self.input_transform = input_transform
def create_kernel_no_ard(self, **kwargs):
return CylindricalKernel(5, MaternKernel(nu=2.5), **kwargs)
