Implementing custom surrogate model in botorch

[mamunm]

If I want to add some custom surrogate model, it needs to have:

1. The num_outputs property
2. A posterior method that returns the multivariate distribution over the queried points.

For 1D problems, I can use GPyTorchPosterior(MultivariateNormal(mu, cov)) where mu is the mean of the predictions and cov is the covariance matrix. It works for 1D problems, but I couldn't get it to work with nD problems.

What I tried:

def posterior(
        self,
        X: Tensor,
        posterior_transform: Optional[Callable[[Posterior], Posterior]] = None,
    ) -> GPyTorchPosterior:
        '''
        Get the posterior.

        Args:
            X (Tensor): Inputs.
            posterior_transform (Callable[[Posterior], Posterior], optional):
                Posterior transform.

        Returns:
            GPyTorchPosterior: Posterior.
        '''
        batch_size = X.shape[0]
        mu, cov = [], []
        for i in range(batch_size):
            temp_mu, temp_cov = self.predict(X[i])
            mu.append(temp_mu.detach().numpy())
            cov.append(temp_cov.detach().numpy())
        mu = torch.tensor(np.concatenate(mu), dtype=X.dtype).reshape(  # type: ignore
            batch_size, self.num_outputs, 1
        )
        cov = torch.tensor(np.concatenate(cov), dtype=X.dtype).reshape(  # type: ignore
            batch_size, self.num_outputs, 1, 1
        )
        return GPyTorchPosterior(MultivariateNormal(mu, cov))

How to combine the mean and covariance of 2 single 1d models to construct the joint distribution for 2d models?

[saitcakmak]

Hi @mamunm. For multi-output posteriors, you may want to use MultitaskMultivariateNormal. See the example here: https://github.com/pytorch/botorch/blob/main/botorch/models/gpytorch.py#L418-L429