Enabling candidate generation with both non_linear_constraints and fixed_features

botorch/generation/gen.py

@@ -114,15 +114,12 @@ def gen_candidates_scipy(
     # if there are fixed features we may optimize over a domain of lower dimension
     reduced_domain = False
     if fixed_features:
-        # TODO: We can support fixed features, see Max's comment on D33551393. We can
-        # consider adding this at a later point.
-        if nonlinear_inequality_constraints:
-            raise NotImplementedError(
-                "Fixed features are not supported when non-linear inequality "
-                "constraints are given."
-            )
-        # if there are no constraints things are straightforward
-        if not (inequality_constraints or equality_constraints):
+        # if there are no constraints, things are straightforward
+        if not (
+            inequality_constraints
+            or equality_constraints
+            or nonlinear_inequality_constraints
+        ):
             reduced_domain = True
         # if there are we need to make sure features are fixed to specific values
         else:
@@ -137,6 +134,7 @@ def gen_candidates_scipy(
             upper_bounds=upper_bounds,
             inequality_constraints=inequality_constraints,
             equality_constraints=equality_constraints,
+            nonlinear_inequality_constraints=nonlinear_inequality_constraints,
         )
         # call the routine with no fixed_features
         clamped_candidates, batch_acquisition = gen_candidates_scipy(
@@ -146,6 +144,7 @@ def gen_candidates_scipy(
             upper_bounds=_no_fixed_features.upper_bounds,
             inequality_constraints=_no_fixed_features.inequality_constraints,
             equality_constraints=_no_fixed_features.equality_constraints,
+            nonlinear_inequality_constraints=_no_fixed_features.nonlinear_inequality_constraints,  # noqa: E501
             options=options,
             fixed_features=None,
             timeout_sec=timeout_sec,
@@ -342,6 +341,7 @@ def gen_candidates_torch(
             upper_bounds=upper_bounds,
             inequality_constraints=None,
             equality_constraints=None,
+            nonlinear_inequality_constraints=None,
         )
 
         # call the routine with no fixed_features

botorch/generation/utils.py

@@ -7,11 +7,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from botorch.acquisition import AcquisitionFunction, FixedFeatureAcquisitionFunction
-from botorch.optim.parameter_constraints import _generate_unfixed_lin_constraints
+from botorch.optim.parameter_constraints import (
+    _generate_unfixed_lin_constraints,
+    _generate_unfixed_nonlin_constraints,
+)
 from torch import Tensor
 
 
@@ -63,6 +66,7 @@ class _NoFixedFeatures:
     upper_bounds: Optional[Union[float, Tensor]]
     inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]]
     equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]]
+    nonlinear_inequality_constraints: Optional[List[Callable]]
 
 
 def _remove_fixed_features_from_optimization(
@@ -73,6 +77,7 @@ def _remove_fixed_features_from_optimization(
     upper_bounds: Optional[Union[float, Tensor]],
     inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
     equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
+    nonlinear_inequality_constraints: Optional[List[Callable]],
 ) -> _NoFixedFeatures:
     """
     Given a set of non-empty fixed features, this function effectively reduces the
@@ -140,11 +145,17 @@ def _remove_fixed_features_from_optimization(
         dimension=d,
         eq=True,
     )
+    nonlinear_inequality_constraints = _generate_unfixed_nonlin_constraints(
+        constraints=nonlinear_inequality_constraints,
+        fixed_features=fixed_features,
+        dimension=d,
+    )
     return _NoFixedFeatures(
         acquisition_function=acquisition_function,
         initial_conditions=initial_conditions,
         lower_bounds=lower_bounds,
         upper_bounds=upper_bounds,
         inequality_constraints=inequality_constraints,
         equality_constraints=equality_constraints,
+        nonlinear_inequality_constraints=nonlinear_inequality_constraints,
     )

botorch/optim/parameter_constraints.py

@@ -312,13 +312,48 @@ def _make_linear_constraints(
     return constraints
 
 
+def _generate_unfixed_nonlin_constraints(
+    constraints: Optional[List[Callable]],
+    fixed_features: Dict[int, float],
+    dimension: int,
+) -> Optional[List[Callable]]:
+    # If constraints is None or an empty list, then return itself
+    if not constraints:
+        return constraints
+
+    selector = []
+    idx_X, idx_f = 0, dimension - len(fixed_features)
+    for i in range(dimension):
+        if i in fixed_features.keys():
+            selector.append(idx_f)
+            idx_f += 1
+        else:
+            selector.append(idx_X)
+            idx_X += 1
+
+    values = torch.tensor([v for v in fixed_features.values()], dtype=torch.double)
+
+    new_constraints = []
+
+    def _wrap_nlc(nlc: Callable) -> Callable:
+        def new_nlc(X: Tensor) -> Tensor:
+            ivalues = values.to(X).expand(*X.shape[:-1], len(fixed_features))
+            X_perm = torch.cat([X, ivalues], dim=-1)
+            return nlc(X_perm[..., selector])
+
+        return new_nlc
+
+    for nlc in constraints:
+        new_constraints.append(_wrap_nlc(nlc=nlc))
+    return new_constraints
+
+
 def _generate_unfixed_lin_constraints(
     constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
     fixed_features: Dict[int, float],
     dimension: int,
     eq: bool,
 ) -> Optional[List[Tuple[Tensor, Tensor, float]]]:
-
     # If constraints is None or an empty list, then return itself
     if not constraints:
         return constraints

test/generation/test_utils.py

@@ -105,6 +105,7 @@ def check_cons(old_cons, new_cons):
                 upper_bounds=upper_bounds,
                 inequality_constraints=inequality_constraints,
                 equality_constraints=equality_constraints,
+                nonlinear_inequality_constraints=None,
             )
             self.assertIsInstance(
                 _no_ff.acquisition_function, FixedFeatureAcquisitionFunction

test/optim/test_optimize.py

@@ -862,6 +862,32 @@ def nlc4(x):
                 torch.allclose(acq_value, torch.tensor(2.45, **tkwargs), atol=1e-3)
             )
 
+            with torch.random.fork_rng():
+                torch.manual_seed(0)
+                batch_initial_conditions = torch.rand(num_restarts, 1, 3, **tkwargs)
+                batch_initial_conditions[..., 0] = 2
+
+            # test with fixed features
+            candidates, acq_value = optimize_acqf(
+                acq_function=mock_acq_function,
+                bounds=bounds,
+                q=1,
+                nonlinear_inequality_constraints=[nlc1, nlc2],
+                batch_initial_conditions=batch_initial_conditions,
+                num_restarts=num_restarts,
+                fixed_features={0: 2},
+            )
+            self.assertEqual(candidates[0, 0], 2.0)
+            self.assertTrue(
+                torch.allclose(
+                    torch.sort(candidates).values,
+                    torch.tensor([[0, 2, 2]], **tkwargs),
+                )
+            )
+            self.assertTrue(
+                torch.allclose(acq_value, torch.tensor(2.8284, **tkwargs), atol=1e-3)
+            )
+
             # Test that an ic_generator object with the same API as
             # gen_batch_initial_conditions returns candidates of the
             # required shape.
@@ -879,22 +905,6 @@ def nlc4(x):
                 )
                 self.assertEqual(candidates.size(), torch.Size([1, 3]))
 
-            # Make sure fixed features aren't supported
-            with self.assertRaisesRegex(
-                NotImplementedError,
-                "Fixed features are not supported when non-linear inequality "
-                "constraints are given.",
-            ):
-                optimize_acqf(
-                    acq_function=mock_acq_function,
-                    bounds=bounds,
-                    q=1,
-                    nonlinear_inequality_constraints=[nlc1, nlc2, nlc3, nlc4],
-                    batch_initial_conditions=batch_initial_conditions,
-                    num_restarts=num_restarts,
-                    fixed_features={0: 0.1},
-                )
-
             # Constraints must be passed in as lists
             with self.assertRaisesRegex(
                 ValueError,

test/optim/test_parameter_constraints.py

@@ -12,6 +12,7 @@
 from botorch.optim.parameter_constraints import (
     _arrayify,
     _generate_unfixed_lin_constraints,
+    _generate_unfixed_nonlin_constraints,
     _make_linear_constraints,
     eval_lin_constraint,
     lin_constraint_jac,
@@ -215,6 +216,42 @@ def test_make_scipy_linear_constraints_unsupported(self):
                 equality_constraints=[(indices, coefficients, 1.0)],
             )
 
+    def test_generate_unfixed_nonlin_constraints(self):
+        def nlc1(x):
+            return 4 - x.sum(dim=-1)
+
+        def nlc2(x):
+            return x[..., 0] - 1
+
+        # first test with one constraint
+        new_constraints = _generate_unfixed_nonlin_constraints(
+            constraints=[nlc1], fixed_features={1: 2.0}, dimension=3
+        )
+        self.assertAllClose(
+            nlc1(torch.tensor([[4.0, 2.0, 2.0]], device=self.device)),
+            new_constraints[0](torch.tensor([[4.0, 2.0]], device=self.device)),
+        )
+        # test with several constraints
+        constraints = [nlc1, nlc2]
+        new_constraints = _generate_unfixed_nonlin_constraints(
+            constraints=constraints, fixed_features={1: 2.0}, dimension=3
+        )
+        for i in range(2):
+            self.assertAllClose(
+                constraints[i](torch.tensor([[4.0, 2.0, 2.0]], device=self.device)),
+                new_constraints[i](torch.tensor([[4.0, 2.0]], device=self.device)),
+            )
+        # test with several constraints and two fixes
+        constraints = [nlc1, nlc2]
+        new_constraints = _generate_unfixed_nonlin_constraints(
+            constraints=constraints, fixed_features={1: 2.0, 2: 1.0}, dimension=3
+        )
+        for i in range(2):
+            self.assertAllClose(
+                constraints[i](torch.tensor([[4.0, 2.0, 1.0]], device=self.device)),
+                new_constraints[i](torch.tensor([[4.0]], device=self.device)),
+            )
+
     def test_generate_unfixed_lin_constraints(self):
         # Case 1: some fixed features are in the indices
         indices = [