Fix double-application of fixed_features to nonlinear constraints (#3332) (#3333)

Summary:
Fixes #3332.

When both `fixed_features` and `nonlinear_inequality_constraints` are passed to `optimize_acqf` / `gen_candidates_scipy`, the fixed features were being applied to the nonlinear constraints **twice**, scrambling the constraint's input and causing spurious SLSQP failures and "infeasible candidate" errors.

## Root cause

In `gen_candidates_scipy`, the nonlinear constraints go through two independent fixed-feature transforms:

1. `_remove_fixed_features_from_optimization` → `_generate_unfixed_nonlin_constraints` wraps each constraint so it accepts the **reduced** (unfixed) variable and re-inserts the fixed features internally (via `cat` + a `selector` permutation).
2. The wrapper passed to `make_scipy_nonlinear_inequality_constraints` was `f_np_wrapper_`, which carries `fixed_features=fixed_features_` and **also** re-inserts the fixed feature (via `fix_features`) before evaluating the constraint.

So the fixed value gets inserted twice. For `d=5`, `fixed_features={0: 1.0}` (`selector=[4,0,1,2,3]`), the reduced `[x1,x2,x3,x4]` first becomes `[1.0,x1,x2,x3,x4]` (from `fix_features`), then the constraint wrapper appends `1.0` and applies `selector`, yielding `[x4, 1.0, x1, x2, x3]` — exactly the scramble reported in #3332.

The acquisition objective wrapper (`f_np_wrapper_`) is correct as-is, since the acquisition function genuinely lives in the full space; only the already-transformed constraints were over-corrected.

## Fix

Use a separate wrapper with `fixed_features=None` for the already-transformed nonlinear constraints, as suggested in the issue.

## Testing

- The existing fixed-features + nonlinear test (`test_optimize_acqf_nonlinear_constraints`) used a permutation-invariant constraint (`4 - x.sum()`) with a `torch.sort()` assertion, so it could not detect coordinate scrambling.
- Added `test_optimize_acqf_nonlinear_constraints_with_fixed_features` with a **coordinate-sensitive** constraint and an order-sensitive assertion. Mutation-checked: it fails on the unpatched code and passes with the fix.
- `test_optimize.py`, `test_gen.py`, `test_parameter_constraints.py`, `test_utils.py` all pass.

Pull Request resolved: #3333

Reviewed By: jandylin

Differential Revision: D109710542

Pulled By: saitcakmak

fbshipit-source-id: bf60b07cfe81b033ab04fd8afb913475f16bf0fe

Author: saitcakmak

botorch/generation/gen.py

@@ -358,9 +358,19 @@ def f(x):
                 nl_ineq_constraints = (
                     _no_fixed_features.nonlinear_inequality_constraints
                 )
+                # The constraints in ``nl_ineq_constraints`` have already been
+                # transformed by ``_generate_unfixed_nonlin_constraints`` to operate
+                # on the reduced (unfixed) variable and re-insert the fixed features
+                # internally. We therefore must NOT re-apply ``fixed_features`` in the
+                # wrapper, otherwise the fixed features would be inserted twice,
+                # producing a wrongly permuted input to the constraint (see #3332).
+                constraint_f_np_wrapper = partial(
+                    f_np_wrapper,
+                    fixed_features=None,
+                )
                 constraints += make_scipy_nonlinear_inequality_constraints(
                     nonlinear_inequality_constraints=nl_ineq_constraints,
-                    f_np_wrapper=f_np_wrapper_,
+                    f_np_wrapper=constraint_f_np_wrapper,
                     x0=x0,
                     shapeX=candidates_.shape,
                 )

test/optim/test_optimize.py

@@ -94,6 +94,17 @@ def forward(self, X):
             return torch.linalg.norm(X, dim=-1).squeeze(-1)
 
 
+class NegSquaredDistanceAcquisitionFunction(AcquisitionFunction):
+    """Negative squared distance to a target, summed over the q-batch."""
+
+    def __init__(self, target: Tensor, model=None):  # noqa: D107
+        super().__init__(model=model)
+        self.register_buffer("target", target)
+
+    def forward(self, X):
+        return -((X - self.target.to(X)) ** 2).sum(dim=-1).sum(dim=-1)
+
+
 class MockOneShotEvaluateAcquisitionFunction(MockOneShotAcquisitionFunction):
     def evaluate(self, X: Tensor, bounds: Tensor):
         return X.sum()
@@ -1337,6 +1348,52 @@ def nlc4(x):
                     raw_samples=16,
                 )
 
+    def test_optimize_acqf_nonlinear_constraints_with_fixed_features(self):
+        # Regression test for
+        # https://github.com/meta-pytorch/botorch/issues/3332
+        # When ``fixed_features`` and ``nonlinear_inequality_constraints`` are
+        # combined, the constraints are transformed by
+        # ``_generate_unfixed_nonlin_constraints`` to operate on the reduced
+        # (unfixed) variable and re-insert the fixed features internally. The
+        # SLSQP wrapper must therefore NOT re-apply ``fixed_features`` on top,
+        # otherwise the fixed features are inserted twice and the constraint
+        # receives a wrongly permuted input.
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            target = torch.tensor([1.0, 0.05, 0.5, 0.5, 0.5], **tkwargs)
+            acqf = NegSquaredDistanceAcquisitionFunction(target=target)
+            bounds = torch.tensor(
+                [[0.0, 0.0, 0.0, 0.0, 0.0], [5.0, 1.0, 1.0, 1.0, 1.0]], **tkwargs
+            )
+
+            # Coordinate-sensitive constraint: x2 * x3 / 2 - x1 >= 0. This is not
+            # permutation invariant, so a scrambled input order would change its
+            # value. Feasible at the target: 0.5 * 0.5 / 2 - 0.05 = 0.075 >= 0.
+            def constraint(x):
+                return x[..., 2] * x[..., 3] / 2.0 - x[..., 1]
+
+            ic = torch.tensor([[[1.0, 0.20, 0.70, 0.70, 0.80]]], **tkwargs)
+            # The initial condition is feasible for the original full constraint.
+            self.assertGreaterEqual(constraint(ic).item(), 0.0)
+
+            candidate, _ = optimize_acqf(
+                acq_function=acqf,
+                bounds=bounds,
+                q=1,
+                num_restarts=1,
+                raw_samples=16,
+                fixed_features={0: 1.0},
+                nonlinear_inequality_constraints=[(constraint, True)],
+                batch_initial_conditions=ic,
+                options={"method": "SLSQP", "batch_limit": 1, "maxiter": 200},
+            )
+            # The fixed feature is respected.
+            self.assertAlmostEqual(candidate[0, 0].item(), 1.0, places=4)
+            # The returned candidate satisfies the original full-dim constraint.
+            self.assertGreaterEqual(constraint(candidate).item(), -1e-6)
+            # And it converges to the (feasible) target.
+            self.assertAllClose(candidate, target.unsqueeze(0), atol=1e-2)
+
     @mock.patch("botorch.optim.optimize.gen_batch_initial_conditions")
     @mock.patch("botorch.optim.optimize.gen_candidates_scipy")
     def test_optimize_acqf_non_linear_constraints_sequential(