Source code for botorch.posteriors.ensemble

#!/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"""
Ensemble posteriors. Used in conjunction with ensemble models.
"""

from __future__ import annotations

from typing import Optional

import torch
from botorch.posteriors.posterior import Posterior
from torch import Tensor


[docs]class EnsemblePosterior(Posterior): r"""Ensemble posterior, that should be used for ensemble models that compute eagerly a finite number of samples per X value as for example a deep ensemble or a random forest.""" def __init__(self, values: Tensor) -> None: r""" Args: values: Values of the samples produced by this posterior as a `(b) x s x q x m` tensor where `m` is the output size of the model and `s` is the ensemble size. """ if values.ndim < 3: raise ValueError("Values has to be at least three-dimensional.") self.values = values @property def ensemble_size(self) -> int: r"""The size of the ensemble""" return self.values.shape[-3] @property def weights(self) -> Tensor: r"""The weights of the individual models in the ensemble. Equally weighted by default.""" return torch.ones(self.ensemble_size) / self.ensemble_size @property def device(self) -> torch.device: r"""The torch device of the posterior.""" return self.values.device @property def dtype(self) -> torch.dtype: r"""The torch dtype of the posterior.""" return self.values.dtype @property def mean(self) -> Tensor: r"""The mean of the posterior as a `(b) x n x m`-dim Tensor.""" return self.values.mean(dim=-3) @property def variance(self) -> Tensor: r"""The variance of the posterior as a `(b) x n x m`-dim Tensor. Computed as the sample variance across the ensemble outputs. """ if self.ensemble_size == 1: return torch.zeros_like(self.values.squeeze(-3)) return self.values.var(dim=-3) def _extended_shape( self, sample_shape: torch.Size = torch.Size() # noqa: B008 ) -> torch.Size: r"""Returns the shape of the samples produced by the posterior with the given `sample_shape`. """ return sample_shape + self.values.shape[:-3] + self.values.shape[-2:]
[docs] def rsample( self, sample_shape: Optional[torch.Size] = None, ) -> Tensor: r"""Sample from the posterior (with gradients). Based on the sample shape, base samples are generated and passed to `rsample_from_base_samples`. Args: sample_shape: A `torch.Size` object specifying the sample shape. To draw `n` samples, set to `torch.Size([n])`. To draw `b` batches of `n` samples each, set to `torch.Size([b, n])`. Returns: Samples from the posterior, a tensor of shape `self._extended_shape(sample_shape=sample_shape)`. """ if sample_shape is None: sample_shape = torch.Size([1]) # get indices as base_samples base_samples = ( torch.multinomial( self.weights, num_samples=sample_shape.numel(), replacement=True, ) .reshape(sample_shape) .to(device=self.device) ) return self.rsample_from_base_samples( sample_shape=sample_shape, base_samples=base_samples )
[docs] def rsample_from_base_samples( self, sample_shape: torch.Size, base_samples: Tensor ) -> Tensor: r"""Sample from the posterior (with gradients) using base samples. This is intended to be used with a sampler that produces the corresponding base samples, and enables acquisition optimization via Sample Average Approximation. Args: sample_shape: A `torch.Size` object specifying the sample shape. To draw `n` samples, set to `torch.Size([n])`. To draw `b` batches of `n` samples each, set to `torch.Size([b, n])`. base_samples: A Tensor of indices as base samples of shape `sample_shape`, typically obtained from `IndexSampler`. This is used for deterministic optimization. The predictions of the ensemble corresponding to the indices are then sampled. Returns: Samples from the posterior, a tensor of shape `self._extended_shape(sample_shape=sample_shape)`. """ if base_samples.shape != sample_shape: raise ValueError("Base samples do not match sample shape.") # move sample axis to front values = self.values.movedim(-3, 0) # sample from the first dimension of values return values[base_samples, ...]