Support equality constraints in random generator polytope sampling (#5178)

Summary:
Pull Request resolved: #5178

Thread equality constraints from the search space through the random adapter
to the `HitAndRunPolytopeSampler`. When rejection sampling falls back to
polytope sampling, equality constraints from both `fixed_features` and
parameter constraints are combined into a single `(C, c)` matrix.

- Update `RandomAdapter._gen` to extract equality constraints via
  `extract_equality_constraints`.
- Add `equality_constraints` parameter to `RandomGenerator.gen()`.
- Add `_combine_equality_constraints` method that merges fixed-feature-based
  and parameter-based equality constraints for the polytope sampler.

Reviewed By: bletham

Differential Revision: D100256488

fbshipit-source-id: c74cc6a82b7e49b2526f9b95cb6c9cf0d10859e3

Author: sdaulton

ax/adapter/random.py

@@ -11,6 +11,7 @@
 
 import numpy as np
 from ax.adapter.adapter_utils import (
+    extract_equality_constraints,
     extract_inequality_constraints,
     extract_search_space_digest,
     get_fixed_features,
@@ -95,6 +96,9 @@ def _gen(
         linear_constraints = extract_inequality_constraints(
             search_space.parameter_constraints, self.parameters
         )
+        equality_constraints_np = extract_equality_constraints(
+            search_space.parameter_constraints, self.parameters
+        )
         # Extract generated points.
         # For normal generators these are used to deduplicate against.
         # For in-sample generators (LILO labeling) they are the selection
@@ -177,6 +181,7 @@ def _gen(
             n=n,
             search_space_digest=search_space_digest,
             linear_constraints=linear_constraints,
+            equality_constraints=equality_constraints_np,
             fixed_features=fixed_features_dict,
             model_gen_options=model_gen_options,
             rounding_func=transform_callback(self.parameters, self.transforms),

ax/adapter/tests/test_random_adapter.py

@@ -109,6 +109,65 @@ def test_gen_w_constraints(self) -> None:
         self.assertEqual(obsf[1].parameters, {"x": 3.0, "y": 4.0, "z": 3.0})
         self.assertTrue(np.array_equal(gen_results.weights, np.array([1.0, 2.0])))
 
+    def test_gen_w_equality_constraints(self) -> None:
+        # Verify that equality constraints from the search space are extracted
+        # and passed through to the generator's gen() call.
+        x = RangeParameter("x", ParameterType.FLOAT, lower=0, upper=1)
+        y = RangeParameter("y", ParameterType.FLOAT, lower=0, upper=1)
+        z = RangeParameter("z", ParameterType.FLOAT, lower=0, upper=1)
+        parameter_constraints = [
+            ParameterConstraint(equality="x + y == 0.5"),
+        ]
+        search_space = SearchSpace([x, y, z], parameter_constraints)
+        experiment = Experiment(search_space=search_space)
+        adapter = RandomAdapter(experiment=experiment, generator=RandomGenerator())
+        with mock.patch.object(
+            adapter.generator,
+            "gen",
+            return_value=(
+                np.array([[0.2, 0.3, 0.4]]),
+                np.array([1.0]),
+            ),
+        ) as mock_gen:
+            adapter._gen(
+                n=1,
+                search_space=search_space,
+                pending_observations={},
+                fixed_features=ObservationFeatures({}),
+                optimization_config=None,
+                model_gen_options=self.model_gen_options,
+            )
+        gen_args = mock_gen.mock_calls[0][2]
+        eq_constraints = gen_args["equality_constraints"]
+        self.assertIsNotNone(eq_constraints)
+        A, b = eq_constraints
+        # x + y = 0.5 => A = [[1, 1, 0]], b = [[0.5]]
+        self.assertTrue(np.array_equal(A, np.array([[1.0, 1.0, 0.0]])))
+        self.assertTrue(np.array_equal(b, np.array([[0.5]])))
+
+    def test_gen_no_equality_constraints(self) -> None:
+        # Verify that equality_constraints is None when there are no equality
+        # constraints on the search space.
+        adapter = RandomAdapter(experiment=self.experiment, generator=RandomGenerator())
+        with mock.patch.object(
+            adapter.generator,
+            "gen",
+            return_value=(
+                np.array([[0.5, 1.5, 2.5]]),
+                np.array([1.0]),
+            ),
+        ) as mock_gen:
+            adapter._gen(
+                n=1,
+                search_space=self.search_space,
+                pending_observations={},
+                fixed_features=ObservationFeatures({}),
+                optimization_config=None,
+                model_gen_options=self.model_gen_options,
+            )
+        gen_args = mock_gen.mock_calls[0][2]
+        self.assertIsNone(gen_args["equality_constraints"])
+
     def test_gen_simple(self) -> None:
         # Test with no constraints, no fixed feature, no pending observations
         search_space = SearchSpace(self.parameters[:2])

ax/generators/random/base.py

@@ -113,6 +113,7 @@ def gen(
         n: int,
         search_space_digest: SearchSpaceDigest,
         linear_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
+        equality_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
         fixed_features: dict[int, float] | None = None,
         model_gen_options: TConfig | None = None,
         rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None,
@@ -127,6 +128,9 @@ def gen(
             linear_constraints: A tuple of (A, b). For k linear constraints on
                 d-dimensional x, A is (k x d) and b is (k x 1) such that
                 A x <= b.
+            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.
             fixed_features: A map {feature_index: value} for features that
                 should be fixed to a particular value during generation.
             model_gen_options: A config dictionary that is passed along to the
@@ -205,9 +209,10 @@ def gen(
                     inequality_constraints=self._convert_inequality_constraints(
                         linear_constraints,
                     ),
-                    equality_constraints=self._convert_equality_constraints(
+                    equality_constraints=self._combine_equality_constraints(
                         d=len(search_space_digest.bounds),
                         fixed_features=fixed_features,
+                        equality_constraints=equality_constraints,
                     ),
                     bounds=self._convert_bounds(bounds=search_space_digest.bounds),
                     interior_point=interior_point,
@@ -353,6 +358,43 @@ def _convert_equality_constraints(
             constraint_matrix[index, fixed_indices[index]] = 1.0
         return constraint_matrix, fixed_vals
 
+    def _combine_equality_constraints(
+        self,
+        d: int,
+        fixed_features: dict[int, float] | None,
+        equality_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
+    ) -> tuple[Tensor, Tensor] | None:
+        """Combine fixed-feature equality constraints with parameter equality
+        constraints into a single (C, c) matrix for the polytope sampler.
+
+        Args:
+            d: Dimension of samples.
+            fixed_features: A map {feature_index: value} for features that
+                should be fixed to a particular value during generation.
+            equality_constraints: A tuple of (A, b) NumPy arrays from
+                ``extract_equality_constraints``.
+
+        Returns:
+            Optional 2-element tuple containing C and c such that Cx = c.
+        """
+        fixed_eq = self._convert_equality_constraints(
+            d=d, fixed_features=fixed_features
+        )
+        param_eq = None
+        if equality_constraints is not None:
+            A = torch.as_tensor(equality_constraints[0], dtype=torch.double)
+            b = torch.as_tensor(equality_constraints[1], dtype=torch.double).squeeze(-1)
+            param_eq = (A, b)
+
+        if fixed_eq is None and param_eq is None:
+            return None
+        if fixed_eq is not None and param_eq is not None:
+            return (
+                torch.cat([fixed_eq[0], param_eq[0]], dim=0),
+                torch.cat([fixed_eq[1], param_eq[1]], dim=0),
+            )
+        return fixed_eq if fixed_eq is not None else param_eq
+
     def _convert_bounds(self, bounds: list[tuple[float, float]]) -> Tensor | None:
         """Helper method to convert bounds list used by the rejectionsampler to the
         tensor format required for the polytope sampler.

ax/generators/random/in_sample.py

@@ -37,6 +37,7 @@ def gen(
         n: int,
         search_space_digest: SearchSpaceDigest,
         linear_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
+        equality_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
         fixed_features: dict[int, float] | None = None,
         model_gen_options: TConfig | None = None,
         rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None,
@@ -49,6 +50,7 @@ def gen(
             search_space_digest: A ``SearchSpaceDigest`` object containing
                 metadata on the features in the datasets.
             linear_constraints: Not used. Accepted for interface compatibility.
+            equality_constraints: Not used. Accepted for interface compatibility.
             fixed_features: Not used. Accepted for interface compatibility.
             model_gen_options: Not used. Accepted for interface compatibility.
             rounding_func: Not used. Accepted for interface compatibility.

ax/generators/random/sobol.py

@@ -81,6 +81,7 @@ def gen(
         n: int,
         search_space_digest: SearchSpaceDigest,
         linear_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
+        equality_constraints: tuple[npt.NDArray, npt.NDArray] | None = None,
         fixed_features: dict[int, float] | None = None,
         model_gen_options: TConfig | None = None,
         rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None,
@@ -95,6 +96,9 @@ def gen(
             linear_constraints: A tuple of (A, b). For k linear constraints on
                 d-dimensional x, A is (k x d) and b is (k x 1) such that
                 A x <= b.
+            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.
             fixed_features: A map {feature_index: value} for features that
                 should be fixed to a particular value during generation.
             rounding_func: A function that rounds an optimization result
@@ -117,6 +121,7 @@ def gen(
             n=n,
             search_space_digest=search_space_digest,
             linear_constraints=linear_constraints,
+            equality_constraints=equality_constraints,
             fixed_features=fixed_features,
             model_gen_options=model_gen_options,
             rounding_func=rounding_func,

ax/generators/tests/test_random.py

@@ -61,6 +61,62 @@ def test_ConvertEqualityConstraints(self) -> None:
         self.assertEqual(C_comparison.any(), True)
         self.assertEqual(self.random_model._convert_equality_constraints(d, None), None)
 
+    def test_CombineEqualityConstraints(self) -> None:
+        d = 4
+        # Both None: returns None.
+        self.assertIsNone(
+            self.random_model._combine_equality_constraints(
+                d=d, fixed_features=None, equality_constraints=None
+            )
+        )
+
+        # Only fixed_features: returns the fixed-feature constraints.
+        fixed_features = {1: 0.5, 3: 0.7}
+        C, c = none_throws(
+            self.random_model._combine_equality_constraints(
+                d=d, fixed_features=fixed_features, equality_constraints=None
+            )
+        )
+        self.assertEqual(C.shape, (2, d))
+        self.assertEqual(c.shape, (2,))
+        self.assertEqual(C[0, 1].item(), 1.0)
+        self.assertEqual(C[1, 3].item(), 1.0)
+        self.assertAlmostEqual(c[0].item(), 0.5)
+        self.assertAlmostEqual(c[1].item(), 0.7)
+
+        # Only equality_constraints: returns the parameter constraints.
+        A_np = np.array([[1.0, 1.0, 0.0, 0.0]])
+        b_np = np.array([[2.0]])
+        C, c = none_throws(
+            self.random_model._combine_equality_constraints(
+                d=d,
+                fixed_features=None,
+                equality_constraints=(A_np, b_np),
+            )
+        )
+        self.assertEqual(C.shape, (1, d))
+        self.assertEqual(c.shape, (1,))
+        self.assertTrue(torch.equal(C, torch.tensor([[1.0, 1.0, 0.0, 0.0]])))
+        self.assertAlmostEqual(c[0].item(), 2.0)
+
+        # Both present: concatenates fixed-feature and parameter constraints.
+        C, c = none_throws(
+            self.random_model._combine_equality_constraints(
+                d=d,
+                fixed_features=fixed_features,
+                equality_constraints=(A_np, b_np),
+            )
+        )
+        # 2 from fixed_features + 1 from equality_constraints = 3 rows.
+        self.assertEqual(C.shape, (3, d))
+        self.assertEqual(c.shape, (3,))
+        # First two rows are from fixed_features (sorted by key: 1, 3).
+        self.assertEqual(C[0, 1].item(), 1.0)
+        self.assertEqual(C[1, 3].item(), 1.0)
+        # Third row is from the parameter equality constraint.
+        self.assertTrue(torch.equal(C[2], torch.tensor([1.0, 1.0, 0.0, 0.0])))
+        self.assertAlmostEqual(c[2].item(), 2.0)
+
     def test_ConvertInequalityConstraints(self) -> None:
         A = np.array([[1, 2], [3, 4]])
         b = np.array([[5], [6]])

ax/generators/tests/test_sobol.py

@@ -9,6 +9,7 @@
 from unittest import mock
 
 import numpy as np
+import torch
 from ax.core.search_space import SearchSpaceDigest
 from ax.exceptions.core import SearchSpaceExhausted
 from ax.generators.random.sobol import SobolGenerator
@@ -244,6 +245,81 @@ def rounding_function(x: npt.NDArray) -> npt.NDArray:
 
         rounding_func.assert_called()
 
+    def test_SobolGeneratorFallbackToPolytopeSamplerWithEqualityConstraints(
+        self,
+    ) -> None:
+        # Ten parameters with sum <= 1 (forces polytope fallback) plus
+        # an equality constraint x0 + x1 = 0.4 to verify equality constraints
+        # are threaded through to the polytope sampler.
+        generator = SobolGenerator(seed=0, fallback_to_sample_polytope=True)
+        ssd = self._create_ssd(n_tunable=10, n_fixed=0)
+        A_ineq = np.ones((1, 10))
+        b_ineq = np.array([1]).reshape((1, 1))
+        # Equality constraint: x0 + x1 = 0.4
+        A_eq = np.zeros((1, 10))
+        A_eq[0, 0] = 1.0
+        A_eq[0, 1] = 1.0
+        b_eq = np.array([[0.4]])
+        generated_points, _ = generator.gen(
+            n=3,
+            search_space_digest=ssd,
+            linear_constraints=(A_ineq, b_ineq),
+            equality_constraints=(A_eq, b_eq),
+            rounding_func=lambda x: x,
+        )
+        self.assertEqual(np.shape(generated_points), (3, 10))
+        # Inequality constraint satisfied.
+        self.assertTrue(np.all(generated_points @ A_ineq.T <= b_ineq + 1e-6))
+        # Equality constraint satisfied: x0 + x1 ≈ 0.4.
+        eq_vals = generated_points @ A_eq.T
+        np.testing.assert_allclose(eq_vals, 0.4, atol=1e-6)
+
+    def test_SobolGeneratorFallbackToPolytopeSamplerWithEqAndFixedFeatures(
+        self,
+    ) -> None:
+        # Verify that equality constraints and fixed features are combined
+        # correctly and passed to the polytope sampler.
+        generator = SobolGenerator(seed=0, fallback_to_sample_polytope=True)
+        ssd = self._create_ssd(n_tunable=10, n_fixed=1)
+        A_ineq = np.insert(np.ones((1, 10)), 10, 0, axis=1)
+        b_ineq = np.array([1]).reshape((1, 1))
+        # Equality constraint: x0 + x1 = 0.3.
+        A_eq = np.zeros((1, 11))
+        A_eq[0, 0] = 1.0
+        A_eq[0, 1] = 1.0
+        b_eq = np.array([[0.3]])
+
+        with mock.patch(
+            "ax.generators.random.base.HitAndRunPolytopeSampler"
+        ) as MockSampler:
+            mock_instance = MockSampler.return_value
+            mock_instance.draw.return_value = torch.rand(1, 11, dtype=torch.double)
+            try:
+                generator.gen(
+                    n=1,
+                    search_space_digest=ssd,
+                    linear_constraints=(A_ineq, b_ineq),
+                    equality_constraints=(A_eq, b_eq),
+                    fixed_features={10: 1},
+                    rounding_func=lambda x: x,
+                )
+            except Exception:
+                pass
+            if MockSampler.called:
+                eq_arg = MockSampler.call_args.kwargs["equality_constraints"]
+                self.assertIsNotNone(eq_arg)
+                C, c = eq_arg
+                # 1 from fixed features + 1 from parameter equality = 2 rows.
+                self.assertEqual(C.shape[0], 2)
+                self.assertEqual(C.shape[1], 11)
+                # Fixed feature: x[10] = 1 (last row from fixed, first in sorted).
+                self.assertEqual(C[0, 10].item(), 1.0)
+                self.assertAlmostEqual(c[0].item(), 1.0)
+                # Parameter equality: x[0] + x[1] = 0.3.
+                self.assertEqual(C[1, 0].item(), 1.0)
+                self.assertEqual(C[1, 1].item(), 1.0)
+                self.assertAlmostEqual(c[1].item(), 0.3)
+
     def test_SobolGeneratorFallbackToPolytopeSamplerWithFixedParam(self) -> None:
         # Ten parameters with sum less than 1. In this example, the rejection
         # sampler gives a search space exhausted error.  Testing fallback to