# Source code for botorch.utils.constraints

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

r"""
Helpers for handling input or outcome constraints.
"""

from __future__ import annotations

from functools import partial
from typing import Callable, List, Optional, Tuple

import torch
from torch import Tensor

[docs]
def get_outcome_constraint_transforms(
outcome_constraints: Optional[Tuple[Tensor, Tensor]]
) -> Optional[List[Callable[[Tensor], Tensor]]]:
r"""Create outcome constraint callables from outcome constraint tensors.

Args:
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is k x m and b is k x 1 such
that A f(x) <= b.

Returns:
A list of callables, each mapping a Tensor of size b x q x m to a
tensor of size b x q, where m is the number of outputs of the model.
Negative values imply feasibility. The callables support broadcasting
(e.g. for calling on a tensor of shape mc_samples x b x q x m).

Example:
>>> # constrain f(x)[0] <= 0
>>> A = torch.tensor([[1., 0.]])
>>> b = torch.tensor([[0.]])
>>> outcome_constraints = get_outcome_constraint_transforms((A, b))
"""
if outcome_constraints is None:
return None
A, b = outcome_constraints

def _oc(a: Tensor, rhs: Tensor, Y: Tensor) -> Tensor:
r"""Evaluate constraints.

Note: einsum multiples Y by a and sums over the m-dimension. Einsum
is ~2x faster than using (Y * a.view(1, 1, -1)).sum(dim-1).

Args:
a: m-dim tensor of weights for the outcomes
rhs: Singleton tensor containing the outcome constraint value
Y: ... x b x q x m tensor of function values

Returns:
A ... x b x q-dim tensor where negative values imply feasibility
"""
lhs = torch.einsum("...m, m", [Y, a])
return lhs - rhs

return [partial(_oc, a, rhs) for a, rhs in zip(A, b)]

[docs]
def get_monotonicity_constraints(
d: int,
descending: bool = False,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
) -> Tuple[Tensor, Tensor]:
"""Returns a system of linear inequalities (A, b) that generically encodes order
constraints on the elements of a d-dimsensional space, i.e. A @ x < b implies
x[i] < x[i + 1] for a d-dimensional vector x.

Idea: Could encode A as sparse matrix, if it is supported well.

Args:
d: Dimensionality of the constraint space, i.e. number of monotonic parameters.
descending: If True, forces the elements of a vector to be monotonically de-
creasing and be monotonically increasing otherwise.
dtype: The dtype of the returned Tensors.
device: The device of the returned Tensors.

Returns:
A tuple of Tensors (A, b) representing the monotonicity constraint as a system
of linear inequalities A @ x < b. A is (d - 1) x d-dimensional and b is
(d - 1) x 1-dimensional.
"""
A = torch.zeros(d - 1, d, dtype=dtype, device=device)
idx = torch.arange(d - 1)
A[idx, idx] = 1
A[idx, idx + 1] = -1
b = torch.zeros(d - 1, 1, dtype=dtype, device=device)
if descending:
A = -A
return A, b

`