Source code for botorch.utils.containers

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

r"""
Containers to standardize inputs into models and acquisition functions.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional

import torch
from botorch.exceptions.errors import UnsupportedError
from torch import Tensor


[docs]@dataclass
class TrainingData:
    r"""Standardized container of model training data for models.

    Properties:
        Xs: A list of tensors, each of shape `batch_shape x n_i x d`,
            where `n_i` is the number of training inputs for the i-th model.
        Ys: A list of tensors, each of shape `batch_shape x n_i x 1`,
            where `n_i` is the number of training observations for the i-th
            (single-output) model.
        Yvars: A list of tensors, each of shape `batch_shape x n_i x 1`,
            where `n_i` is the number of training observations of the
            observation noise for the i-th  (single-output) model.
            If `None`, the observation noise level is unobserved.
    """

    Xs: List[Tensor]  # `batch_shape x n_i x 1`
    Ys: List[Tensor]  # `batch_shape x n_i x 1`
    Yvars: Optional[List[Tensor]] = None  # `batch_shape x n_i x 1`

    def __post_init__(self):
        self._is_block_design = all(torch.equal(X, self.Xs[0]) for X in self.Xs[1:])

[docs]    @classmethod
    def from_block_design(cls, X: Tensor, Y: Tensor, Yvar: Optional[Tensor] = None):
        r"""Construct a TrainingData object from a block design description.

        Args:
            X: A `batch_shape x n x d` tensor of training points (shared across
                all outcomes).
            Y: A `batch_shape x n x m` tensor of training observations.
            Yvar: A `batch_shape x n x m` tensor of training noise variance
                observations, or `None`.

        Returns:
            The `TrainingData` object (with `is_block_design=True`).
        """
        return cls(
            Xs=[X for _ in range(Y.shape[-1])],
            Ys=list(torch.split(Y, 1, dim=-1)),
            Yvars=None if Yvar is None else list(torch.split(Yvar, 1, dim=-1)),
        )

    @property
    def is_block_design(self) -> bool:
        r"""Indicates whether training data is a "block design".

        Block designs are designs in which all outcomes are observed
        at the same training inputs.
        """
        return self._is_block_design

    @property
    def X(self) -> Tensor:
        r"""The training inputs (block-design only).

        This raises an `UnsupportedError` in the non-block-design case.
        """
        if not self.is_block_design:
            raise UnsupportedError
        return self.Xs[0]

    @property
    def Y(self) -> Tensor:
        r"""The training observations (block-design only).

        This raises an `UnsupportedError` in the non-block-design case.
        """
        if not self.is_block_design:
            raise UnsupportedError
        return torch.cat(self.Ys, dim=-1)

    @property
    def Yvar(self) -> Optional[List[Tensor]]:
        r"""The training observations's noise variance (block-design only).

        This raises an `UnsupportedError` in the non-block-design case.
        """
        if self.Yvars is not None:
            if not self.is_block_design:
                raise UnsupportedError
            return torch.cat(self.Yvars, dim=-1)

    def __eq__(self, other: TrainingData) -> bool:
        # Check for `None` Yvars and unequal attribute lengths upfront.
        if self.Yvars is None or other.Yvars is None:
            if not (self.Yvars is other.Yvars is None):
                return False
        else:
            if len(self.Yvars) != len(other.Yvars):
                return False

        if len(self.Xs) != len(other.Xs) or len(self.Ys) != len(other.Ys):
            return False

        return (  # Deep-check equality of attributes.
            all(
                torch.equal(self_X, other_X)
                for self_X, other_X in zip(self.Xs, other.Xs)
            )
            and all(
                torch.equal(self_Y, other_Y)
                for self_Y, other_Y in zip(self.Ys, other.Ys)
            )
            and (
                self.Yvars is other.Yvars is None
                or all(
                    torch.equal(self_Yvar, other_Yvar)
                    for self_Yvar, other_Yvar in zip(self.Yvars, other.Yvars)
                )
            )
        )