|
11 | 11 | from __future__ import annotations
|
12 | 12 |
|
13 | 13 | from collections.abc import Callable
|
14 |
| - |
15 | 14 | from functools import partial
|
16 | 15 | from typing import Union
|
17 | 16 |
|
|
26 | 25 | str, Union[str, Callable[[np.ndarray], float], Callable[[np.ndarray], np.ndarray]]
|
27 | 26 | ]
|
28 | 27 | NLC_TOL = -1e-6
|
| 28 | +INTRA_POINT_CONST_ERR: str = "Only intra-point constraints are supported." |
29 | 29 |
|
30 | 30 |
|
31 | 31 | def make_scipy_bounds(
|
@@ -601,3 +601,58 @@ def make_scipy_nonlinear_inequality_constraints(
|
601 | 601 | shapeX=shapeX,
|
602 | 602 | )
|
603 | 603 | return scipy_nonlinear_inequality_constraints
|
| 604 | + |
| 605 | + |
| 606 | +def evaluate_feasibility( |
| 607 | + X: Tensor, |
| 608 | + inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None, |
| 609 | + equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None, |
| 610 | + nonlinear_inequality_constraints: list[tuple[Callable, bool]] | None = None, |
| 611 | +) -> Tensor: |
| 612 | + r"""Evaluate feasibility of a set of points. Only supports intra-point constraints. |
| 613 | +
|
| 614 | + Args: |
| 615 | + X: A tensor of points of shape `batch_shape x d`. |
| 616 | + inequality_constraints: A list of tuples (indices, coefficients, rhs), |
| 617 | + with each tuple encoding an inequality constraint of the form |
| 618 | + `\sum_i (X[indices[i]] * coefficients[i]) >= rhs`. `indices` and |
| 619 | + `coefficients` should be torch tensors. See the docstring of |
| 620 | + `make_scipy_linear_constraints` for an example. |
| 621 | + equality_constraints: A list of tuples (indices, coefficients, rhs), |
| 622 | + with each tuple encoding an equality constraint of the form |
| 623 | + `\sum_i (X[indices[i]] * coefficients[i]) = rhs`. See the docstring of |
| 624 | + `make_scipy_linear_constraints` for an example. |
| 625 | + nonlinear_inequality_constraints: A list of tuples representing the nonlinear |
| 626 | + inequality constraints. The first element in the tuple is a callable |
| 627 | + representing a constraint of the form `callable(x) >= 0`. The `callable()` |
| 628 | + takes in an one-dimensional tensor of shape `d` and returns a scalar. The |
| 629 | + second element is a boolean, indicating if it is an intra-point or |
| 630 | + inter-point constraint (`True` for intra-point. `False` for |
| 631 | + inter-point). Only `True` is supported here. For more information on |
| 632 | + intra-point vs inter-point constraints, see the docstring of the |
| 633 | + `inequality_constraints` argument to `optimize_acqf()`. |
| 634 | +
|
| 635 | + Returns: |
| 636 | + A boolean tensor of shape `batch` denoting whether each point is feasible. |
| 637 | + """ |
| 638 | + is_feasible = torch.ones(X.shape[:-1], device=X.device, dtype=torch.bool) |
| 639 | + if inequality_constraints is not None: |
| 640 | + for idx, coef, rhs in inequality_constraints: |
| 641 | + if idx.ndim != 1: |
| 642 | + raise UnsupportedError(INTRA_POINT_CONST_ERR) |
| 643 | + is_feasible &= (X[..., idx] * coef).sum(dim=-1) >= rhs |
| 644 | + if equality_constraints is not None: |
| 645 | + for idx, coef, rhs in equality_constraints: |
| 646 | + if idx.ndim != 1: |
| 647 | + raise UnsupportedError(INTRA_POINT_CONST_ERR) |
| 648 | + is_feasible &= (X[..., idx] * coef).sum(dim=-1) == rhs |
| 649 | + if nonlinear_inequality_constraints is not None: |
| 650 | + for const, intra in nonlinear_inequality_constraints: |
| 651 | + if not intra: |
| 652 | + raise UnsupportedError(INTRA_POINT_CONST_ERR) |
| 653 | + is_feasible &= torch.tensor( |
| 654 | + [const(x) >= NLC_TOL for x in X.view(-1, X.shape[-1])], |
| 655 | + device=X.device, |
| 656 | + dtype=torch.bool, |
| 657 | + ).view_as(is_feasible) |
| 658 | + return is_feasible |
0 commit comments