avidereta / risk-averse-hetero-bo Goto Github PK

class qRiskAverseUpperConfidenceBound(MCAcquisitionFunction):

    def __init__(
            self,
            model: Model,
            model_varproxi: Model,
            beta: Union[float, Tensor],
            beta_varproxi: Union[float, Tensor],
            gamma: Union[float, Tensor],
            objective: Optional[ScalarizedObjective] = None,
            maximize: bool = True,
            posterior_transform: Optional[PosteriorTransform] = None,
            X_pending: Optional[Tensor] = None,
    ) -> None:
        super().__init__(
            model=model,
            objective=objective,
            posterior_transform=posterior_transform,
            X_pending=X_pending)
        
        self.maximize = maximize
        self.model_varproxi = model_varproxi
        if not torch.is_tensor(beta):
            beta = torch.tensor([beta])
        self.register_buffer("beta", beta)
        if not torch.is_tensor(beta_varproxi):
            beta_varproxi = torch.tensor([beta_varproxi])
        #         self.beta_varproxi = beta_varproxi
        self.register_buffer("beta_varproxi", beta_varproxi)
        if not torch.is_tensor(gamma):
            gamma = torch.tensor(gamma)
        #         self.gamma = gamma
        self.register_buffer("gamma", gamma)

    def _get_posterior_varproxi(self, X: Tensor, check_single_output: bool = True) -> Posterior:
        self.model_varproxi.eval()
        posterior = self.model_varproxi(X)
        # if check_single_output:
        #     if posterior.event_shape[-1] != 1:
        #         raise UnsupportedError(
        #             "Multi-Output posteriors are not supported for acquisition "
        #             f"functions of type {self.__class__.__name__}"
        #         )
        return posterior
    
    def _get_posterior(self,X):
        self.model.eval()
        posterior = self.model(X)
        return posterior

    @concatenate_pending_points
    @t_batch_mode_transform()
    def forward(self, X: Tensor) -> Tensor:

        self.beta = self.beta.to(X)
        posterior = self._get_posterior(X=X)
        # batch_shape = X.shape[:-2]
        # mean = posterior.mean.view(batch_shape)
        # variance = posterior.variance.view(batch_shape)
        mean = posterior.mean
        variance = posterior.variance
        delta = (self.beta.expand_as(variance) * variance).sqrt()

        self.beta_varproxi = self.beta_varproxi.to(X)
        posterior_varproxi = self._get_posterior_varproxi(X=X)
        # mean_varproxi = posterior_varproxi.mean.view(batch_shape)
        # variance_varproxi = posterior_varproxi.variance.view(batch_shape)
        mean_varproxi = posterior_varproxi.mean
        variance_varproxi = posterior_varproxi.variance
        delta_varproxi = (self.beta_varproxi.expand_as(variance_varproxi) * variance_varproxi).sqrt()

        # ucb = ucb_f - gamma*lcb_{rho}
        if self.maximize:
            value =  (mean + delta - self.gamma.expand_as(mean_varproxi) * (mean_varproxi - delta_varproxi))
        # lcb = lcb_f - gamma*ucb_{rho}
        else:
            value =  (mean - delta - self.gamma.expand_as(mean_varproxi) * (mean_varproxi + delta_varproxi))

        return value.mean(dim=1)