# Source code for botorch.acquisition.joint_entropy_search

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# LICENSE file in the root directory of this source tree.

r"""
Acquisition function for joint entropy search (JES). The code utilizes the
implementation designed for the multi-objective batch setting.

"""

from __future__ import annotations

from typing import Any, Optional

from botorch.acquisition.multi_objective.joint_entropy_search import (
qLowerBoundMultiObjectiveJointEntropySearch,
)
from botorch.acquisition.multi_objective.utils import compute_sample_box_decomposition
from botorch.models.model import Model
from botorch.utils.transforms import concatenate_pending_points, t_batch_mode_transform
from torch import Tensor

[docs]class qLowerBoundJointEntropySearch(qLowerBoundMultiObjectiveJointEntropySearch):
r"""The acquisition function for the Joint Entropy Search, where the batches
q > 1 are supported through the lower bound formulation.

This acquisition function computes the mutual information between the observation
at a candidate point X and the optimal input-output pair.

See [Tu2022]_ for a discussion on the estimation procedure.

NOTES:
(i) The estimated acquisition value could be negative.

(ii) The lower bound batch acquisition function might not be monotone in the
sense that adding more elements to the batch does not necessarily increase the
acquisition value. Specifically, the acquisition value can become smaller when
"""

def __init__(
self,
model: Model,
optimal_inputs: Tensor,
optimal_outputs: Tensor,
maximize: bool = True,
hypercell_bounds: Tensor = None,
X_pending: Optional[Tensor] = None,
estimation_type: str = "LB",
num_samples: int = 64,
**kwargs: Any,
) -> None:
r"""Joint entropy search acquisition function.

Args:
model: A fitted single-outcome model.
optimal_inputs: A num_samples x d-dim tensor containing the sampled
optimal inputs of dimension d. We assume for simplicity that each
sample only contains one optimal set of inputs.
optimal_outputs: A num_samples x 1-dim Tensor containing the optimal
set of objectives of dimension 1.
maximize: If true, we consider a maximization problem.
hypercell_bounds:  A num_samples x 2 x J x 1-dim Tensor containing the
hyper-rectangle bounds for integration, where J is the number of
hyper-rectangles. By default, the problem is assumed to be
unconstrained and therefore the region of integration for a sample
(x*, y*) is a J=1 hyper-rectangle of the form  (-infty, y^*]
for a maximization problem and [y^*, +infty) for a minimization
problem. In the constrained setting, the region of integration also
includes the infeasible space.
X_pending: A m x d-dim Tensor of m design points that have been
submitted for function evaluation, but have not yet been evaluated.
estimation_type: A string to determine which entropy estimate is
computed: "0", "LB", "LB2", or "MC". In the single-objective
setting, "LB" is equivalent to "LB2".
num_samples: The number of Monte Carlo samples used for the Monte Carlo
estimate.
"""
if hypercell_bounds is None:
hypercell_bounds = compute_sample_box_decomposition(
pareto_fronts=optimal_outputs.unsqueeze(-2), maximize=maximize
)

super().__init__(
model=model,
pareto_sets=optimal_inputs.unsqueeze(-2),
pareto_fronts=optimal_outputs.unsqueeze(-2),
hypercell_bounds=hypercell_bounds,
X_pending=X_pending,
estimation_type=estimation_type,
num_samples=num_samples,
)

[docs]    @concatenate_pending_points
@t_batch_mode_transform()
def forward(self, X: Tensor) -> Tensor:
r"""Evaluates qLowerBoundJointEntropySearch at the design points X.

Args:
X: A batch_shape x q x d-dim Tensor of batch_shape t-batches with q
d-dim design points each.

Returns:
A batch_shape-dim Tensor of acquisition values at the given design
points X.
"""

return self._compute_lower_bound_information_gain(X)