Source code for botorch.utils.low_rank

#!/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

import torch
from botorch.exceptions.errors import BotorchError
from botorch.posteriors.base_samples import _reshape_base_samples_non_interleaved
from botorch.posteriors.gpytorch import GPyTorchPosterior
from gpytorch.distributions.multitask_multivariate_normal import (
    MultitaskMultivariateNormal,
)
from linear_operator.operators import BlockDiagLinearOperator, LinearOperator

from linear_operator.utils.cholesky import psd_safe_cholesky
from linear_operator.utils.errors import NanError
from torch import Tensor


[docs]def extract_batch_covar(mt_mvn: MultitaskMultivariateNormal) -> LinearOperator:
    r"""Extract a batched independent covariance matrix from an MTMVN.

    Args:
        mt_mvn: A multi-task multivariate normal with a block diagonal
            covariance matrix.

    Returns:
        A lazy covariance matrix consisting of a batch of the blocks of
            the diagonal of the MultitaskMultivariateNormal.

    """
    lazy_covar = mt_mvn.lazy_covariance_matrix
    if not isinstance(lazy_covar, BlockDiagLinearOperator):
        raise BotorchError(
            f"Expected BlockDiagLinearOperator, but got {type(lazy_covar)}."
        )
    return lazy_covar.base_linear_op


def _reshape_base_samples(
    base_samples: Tensor, sample_shape: torch.Size, posterior: GPyTorchPosterior
) -> Tensor:
    r"""Manipulate shape of base_samples to match `MultivariateNormal.rsample`.

    This ensure that base_samples are used in the same way as in
    gpytorch.distributions.MultivariateNormal. For CBD, it is important to ensure
    that the same base samples are used for the in-sample points here and in the
    cached box decompositions.

    Args:
        base_samples: The base samples.
        sample_shape: The sample shape.
        posterior: The joint posterior is over (X_baseline, X).

    Returns:
        Reshaped and expanded base samples.
    """
    mvn = posterior.distribution
    loc = mvn.loc
    peshape = posterior._extended_shape()
    base_samples = base_samples.view(
        sample_shape + torch.Size([1] * (loc.ndim - 1)) + peshape[-2:]
    ).expand(sample_shape + loc.shape[:-1] + peshape[-2:])
    if posterior._is_mt:
        base_samples = _reshape_base_samples_non_interleaved(
            mvn=posterior.distribution,
            base_samples=base_samples,
            sample_shape=sample_shape,
        )
    base_samples = base_samples.reshape(
        -1, *loc.shape[:-1], mvn.lazy_covariance_matrix.shape[-1]
    )
    base_samples = base_samples.permute(*range(1, loc.dim() + 1), 0)
    return base_samples.reshape(
        *peshape[:-2],
        peshape[-1],
        peshape[-2],
        *sample_shape,
    )


[docs]def sample_cached_cholesky(
    posterior: GPyTorchPosterior,
    baseline_L: Tensor,
    q: int,
    base_samples: Tensor,
    sample_shape: torch.Size,
    max_tries: int = 6,
) -> Tensor:
    r"""Get posterior samples at the `q` new points from the joint multi-output
    posterior.

    Args:
        posterior: The joint posterior is over (X_baseline, X).
        baseline_L: The baseline lower triangular cholesky factor.
        q: The number of new points in X.
        base_samples: The base samples.
        sample_shape: The sample shape.
        max_tries: The number of tries for computing the Cholesky
            decomposition with increasing jitter.


    Returns:
        A `sample_shape x batch_shape x q x m`-dim tensor of posterior
            samples at the new points.
    """
    # compute bottom left covariance block
    mvn = posterior.distribution
    lazy_covar = (
        extract_batch_covar(mt_mvn=mvn)
        if isinstance(mvn, MultitaskMultivariateNormal)
        else mvn.lazy_covariance_matrix
    )
    # Get the `q` new rows of the batched covariance matrix
    bottom_rows = lazy_covar[..., -q:, :].to_dense()
    # The covariance in block form is:
    # [K(X_baseline, X_baseline), K(X_baseline, X)]
    # [K(X, X_baseline), K(X, X)]
    # bl := K(X, X_baseline)
    # br := K(X, X)
    # Get bottom right block of new covariance
    bl, br = bottom_rows.split([bottom_rows.shape[-1] - q, q], dim=-1)
    # Solve Ax = b
    # where A = K(X_baseline, X_baseline) and b = K(X, X_baseline)^T
    # and bl_chol := x^T
    # bl_chol is the new `(batch_shape) x q x n`-dim bottom left block
    # of the cholesky decomposition
    bl_chol = torch.linalg.solve_triangular(
        baseline_L, bl.transpose(-2, -1), upper=False
    ).transpose(-2, -1)

    # Compute the new bottom right block of the Cholesky
    # decomposition via:
    # Cholesky(K(X, X) - bl_chol @ bl_chol^T)
    br_to_chol = br - bl_chol @ bl_chol.transpose(-2, -1)
    # TODO: technically we should make sure that we add a
    # consistent nugget to the cached covariance and the new block
    br_chol = psd_safe_cholesky(br_to_chol, max_tries=max_tries)
    # Create a `(batch_shape) x q x (n+q)`-dim tensor containing the
    # `q` new bottom rows of the Cholesky decomposition
    new_Lq = torch.cat([bl_chol, br_chol], dim=-1)
    mean = posterior.distribution.mean
    base_samples = _reshape_base_samples(
        base_samples=base_samples,
        sample_shape=sample_shape,
        posterior=posterior,
    )
    if not isinstance(posterior.distribution, MultitaskMultivariateNormal):
        # add output dim
        mean = mean.unsqueeze(-1)
        # add batch dim corresponding to output dim
        new_Lq = new_Lq.unsqueeze(-3)
    new_mean = mean[..., -q:, :]
    res = (
        new_Lq.matmul(base_samples)
        .add(new_mean.transpose(-1, -2).unsqueeze(-1))
        .permute(-1, *range(posterior.distribution.loc.dim() - 1), -2, -3)
        .contiguous()
    )
    contains_nans = torch.isnan(res).any()
    contains_infs = torch.isinf(res).any()
    if contains_nans or contains_infs:
        suffix_args = []
        if contains_nans:
            suffix_args.append("nans")
        if contains_infs:
            suffix_args.append("infs")
        suffix = " and ".join(suffix_args)
        raise NanError(f"Samples contain {suffix}.")
    return res