Source code for botorch.models.kernels.downsampling

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import annotations

from typing import Optional

import torch
from gpytorch.constraints import Interval, Positive
from gpytorch.kernels import Kernel
from gpytorch.priors import Prior
from torch import Tensor

[docs] class DownsamplingKernel(Kernel): r"""GPyTorch Downsampling Kernel. Computes a covariance matrix based on the down sampling kernel between inputs `x_1` and `x_2` (we expect `d = 1`): K(\mathbf{x_1}, \mathbf{x_2}) = c + (1 - x_1)^(1 + delta) * (1 - x_2)^(1 + delta). where `c` is an offset parameter, and `delta` is a power parameter. """ def __init__( self, power_prior: Optional[Prior] = None, offset_prior: Optional[Prior] = None, power_constraint: Optional[Interval] = None, offset_constraint: Optional[Interval] = None, **kwargs, ): r""" Args: power_constraint: Constraint to place on power parameter. Default is `Positive`. power_prior: Prior over the power parameter. offset_constraint: Constraint to place on offset parameter. Default is `Positive`. active_dims: List of data dimensions to operate on. `len(active_dims)` should equal `num_dimensions`. """ super().__init__(**kwargs) if power_constraint is None: power_constraint = Positive() if offset_constraint is None: offset_constraint = Positive() self.register_parameter( name="raw_power", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)), ) self.register_parameter( name="raw_offset", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)), ) if power_prior is not None: self.register_prior( "power_prior", power_prior, lambda m: m.power, lambda m, v: m._set_power(v), ) self.register_constraint("raw_power", power_constraint) if offset_prior is not None: self.register_prior( "offset_prior", offset_prior, lambda m: m.offset, lambda m, v: m._set_offset(v), ) self.register_constraint("raw_offset", offset_constraint) @property def power(self) -> Tensor: return self.raw_power_constraint.transform(self.raw_power) @power.setter def power(self, value: Tensor) -> None: self._set_power(value) def _set_power(self, value: Tensor) -> None: if not torch.is_tensor(value): value = torch.as_tensor(value).to(self.raw_power) self.initialize(raw_power=self.raw_power_constraint.inverse_transform(value)) @property def offset(self) -> Tensor: return self.raw_offset_constraint.transform(self.raw_offset) @offset.setter def offset(self, value: Tensor) -> None: self._set_offset(value) def _set_offset(self, value: Tensor) -> None: if not torch.is_tensor(value): value = torch.as_tensor(value).to(self.raw_offset) self.initialize(raw_offset=self.raw_offset_constraint.inverse_transform(value)) def forward( self, x1: Tensor, x2: Tensor, diag: Optional[bool] = False, last_dim_is_batch: Optional[bool] = False, **params, ) -> Tensor: offset = self.offset exponent = 1 + self.power if last_dim_is_batch: x1 = x1.transpose(-1, -2).unsqueeze(-1) x2 = x2.transpose(-1, -2).unsqueeze(-1) x1_ = 1 - x1 x2_ = 1 - x2 if diag: return offset + (x1_ * x2_).sum(dim=-1).pow(exponent) offset = offset.unsqueeze(-1) # unsqueeze enables batch evaluation exponent = exponent.unsqueeze(-1) # unsqueeze enables batch evaluation return offset + x1_.pow(exponent) @ x2_.transpose(-2, -1).pow(exponent)