Enforce equality constraints in random generator rejection sampling (#5182)

Summary:
Pull Request resolved: #5182

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485

fbshipit-source-id: 540e45ad2fd4852ecff7dc7500930758d4fed979

Author: sdaulton

ax/generators/random/base.py

@@ -154,6 +154,13 @@ def gen(
             max_draws = model_gen_options.get("max_rs_draws", DEFAULT_MAX_RS_DRAWS)
             max_draws = int(assert_is_instance_of_tuple(max_draws, (int, float)))
         try:
+            # With equality constraints, unconstrained sampling has probability
+            # zero of producing feasible points, so skip straight to polytope
+            # sampling.
+            if equality_constraints is not None:
+                raise SearchSpaceExhausted(
+                    "Equality constraints require polytope sampling."
+                )
             # Always rejection sample, but this only rejects if there are
             # constraints or actual duplicates and deduplicate is specified.
             # If rejection sampling fails, fall back to polytope sampling.
@@ -184,11 +191,15 @@ def gen(
                 num_generated = (
                     len(generated_points) if generated_points is not None else 0
                 )
-                interior_point = (  # A feasible point of shape `d x 1`.
-                    torch.from_numpy(generated_points[-1].reshape((-1, 1))).double()
-                    if generated_points is not None
-                    else None
-                )
+                # Use a previously generated point as the interior point
+                # hint, but only if it's likely feasible. When equality
+                # constraints are present, previous points (generated
+                # without those constraints) won't satisfy them.
+                interior_point: torch.Tensor | None = None
+                if generated_points is not None and equality_constraints is None:
+                    interior_point = torch.from_numpy(
+                        generated_points[-1].reshape((-1, 1))
+                    ).double()
                 kwargs = {"n_burnin": 100, "n_thinning": 20}
                 kwargs.update(self.polytope_sampler_kwargs)
                 polytope_sampler: HitAndRunPolytopeSampler = HitAndRunPolytopeSampler(
@@ -229,6 +240,7 @@ def gen_polytope_sampler(
                     fixed_features=fixed_features,
                     rounding_func=rounding_func,
                     existing_points=generated_points,
+                    equality_constraints=equality_constraints,
                 )
             else:
                 raise e

ax/generators/utils.py

@@ -68,6 +68,7 @@ def rejection_sample(
     fixed_features: dict[int, float] | None = None,
     rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None,
     existing_points: npt.NDArray | None = None,
+    equality_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
 ) -> tuple[npt.NDArray, int]:
     """Rejection sample in parameter space.
 
@@ -96,6 +97,9 @@ def rejection_sample(
         existing_points: A set of previously generated points to use
             for deduplication. These should be provided in the parameter
             space model operates in.
+        equality_constraints: A tuple of (A, b). For k equality constraints
+            on d-dimensional x, A is (k x d) and b is (k x 1) such that
+            A x = b.
 
     Returns:
         2-element tuple containing the generated points and the number of
@@ -124,9 +128,26 @@ def rejection_sample(
         )[0]
 
         # Check parameter constraints, always in raw transformed space.
+        has_constraints = (
+            linear_constraints is not None or equality_constraints is not None
+        )
         if linear_constraints is not None:
             all_constraints_satisfied, _ = check_param_constraints(
-                linear_constraints=linear_constraints, point=point
+                linear_constraints=linear_constraints,
+                point=point,
+                equality_constraints=equality_constraints,
+            )
+        elif equality_constraints is not None:
+            # No inequality constraints but have equality constraints.
+            # Use a dummy (0, d) inequality matrix so check_param_constraints works.
+            dummy_ineq = (
+                np.zeros((0, len(point))),
+                np.zeros((0, 1)),
+            )
+            all_constraints_satisfied, _ = check_param_constraints(
+                linear_constraints=dummy_ineq,
+                point=point,
+                equality_constraints=equality_constraints,
             )
         else:
             all_constraints_satisfied = True
@@ -140,9 +161,15 @@ def rejection_sample(
                 # Re-check constraints after rounding for discrete parameters
                 # (e.g. numerical choice parameters) because rounding can push values
                 # in a direction that violates sum constraints.
-                if linear_constraints is not None:
+                if has_constraints:
+                    ineq = linear_constraints or (
+                        np.zeros((0, len(point))),
+                        np.zeros((0, 1)),
+                    )
                     all_constraints_satisfied, _ = check_param_constraints(
-                        linear_constraints=linear_constraints, point=point
+                        linear_constraints=ineq,
+                        point=point,
+                        equality_constraints=equality_constraints,
                     )
                     if not all_constraints_satisfied:
                         attempted_draws += 1
@@ -228,6 +255,7 @@ def add_fixed_features(
 def check_param_constraints(
     linear_constraints: tuple[npt.NDArray, npt.NDArray],
     point: npt.NDArray,
+    equality_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
 ) -> tuple[bool, npt.NDArray]:
     """Check if a point satisfies parameter constraints.
 
@@ -236,6 +264,9 @@ def check_param_constraints(
             d-dimensional x, A is (k x d) and b is (k x 1) such that
             A x <= b.
         point: A candidate point in d-dimensional space, as a (1 x d) matrix.
+        equality_constraints: A tuple of (A, b). For k equality constraints on
+            d-dimensional x, A is (k x d) and b is (k x 1) such that
+            A x = b.
 
     Returns:
         2-element tuple containing
@@ -246,6 +277,17 @@ def check_param_constraints(
     constraints_satisfied = (
         linear_constraints[0] @ np.expand_dims(point, axis=1) <= linear_constraints[1]
     )
+    if equality_constraints is not None:
+        eq_satisfied = (
+            np.abs(
+                equality_constraints[0] @ np.expand_dims(point, axis=1)
+                - equality_constraints[1]
+            )
+            <= 1e-8
+        )
+        constraints_satisfied = np.concatenate(
+            [constraints_satisfied, eq_satisfied], axis=0
+        )
     if np.all(constraints_satisfied):
         return True, np.array([])
     else: