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)