add support for general box bounds to RandomGenerators. (facebook#4932)

Summary:
Pull Request resolved: facebook#4932

Currently RandomGenerators require that the space be mapped to [0, 1]^d, for the convenience of random number generators operating on that space by default. In a stacked diff I would like to super() gen on a random generator but need the returned points to be in raw space, no UnitX transformation.

It is a relatively small change for RandomGenerators to be capable of operating with arbitrary bounds, which change this diff makes.

Reviewed By: saitcakmak, eonofrey

Differential Revision: D94118103

fbshipit-source-id: 7d604fbc89eb01ba2e503f6194cf5474b2f4c21b

Author: bletham

ax/generators/random/base.py

@@ -23,7 +23,6 @@
     DEFAULT_MAX_RS_DRAWS,
     rejection_sample,
     tunable_feature_indices,
-    validate_bounds,
 )
 from ax.utils.common.docutils import copy_doc
 from ax.utils.common.logger import get_logger
@@ -85,6 +84,7 @@ def __init__(
         # Used for deduplication.
         self.fallback_to_sample_polytope = fallback_to_sample_polytope
         self.polytope_sampler_kwargs: dict[str, Any] = polytope_sampler_kwargs or {}
+        self._bounds: npt.NDArray = np.empty((0, 2))
         self.attempted_draws: int = 0
         if generated_points is not None:
             # generated_points was deprecated in Ax 1.0.0, so it can now be reaped.
@@ -146,15 +146,8 @@ def gen(
         tf_indices = tunable_feature_indices(
             bounds=search_space_digest.bounds, fixed_features=fixed_features
         )
-        if fixed_features:
-            fixed_feature_indices = np.array(list(fixed_features.keys()))
-        else:
-            fixed_feature_indices = np.array([])
+        self._bounds = np.array(search_space_digest.bounds)  # (d, 2)
 
-        validate_bounds(
-            bounds=search_space_digest.bounds,
-            fixed_feature_indices=fixed_feature_indices,
-        )
         max_draws = DEFAULT_MAX_RS_DRAWS
         discrete_indices = set(search_space_digest.discrete_choices.keys())
         continuous_indices = {
@@ -282,22 +275,29 @@ def _gen_unconstrained(
             An (n x d) array of generated points.
 
         """
-        tunable_points = self._gen_samples(n=n, tunable_d=len(tunable_feature_indices))
+        tunable_bounds = self._bounds[tunable_feature_indices]
+        tunable_points = self._gen_samples(
+            n=n,
+            tunable_d=len(tunable_feature_indices),
+            bounds=tunable_bounds,
+        )
         return add_fixed_features(
             tunable_points=tunable_points,
             d=d,
             tunable_feature_indices=tunable_feature_indices,
             fixed_features=fixed_features,
         )
 
-    def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray:
-        """Generate n samples on [0, 1]^d.
+    def _gen_samples(self, n: int, tunable_d: int, bounds: npt.NDArray) -> npt.NDArray:
+        """Generate n samples within the given bounds.
 
         Args:
             n: Number of points to generate.
+            tunable_d: Number of tunable dimensions.
+            bounds: A (tunable_d x 2) array of (lower, upper) bounds.
 
         Returns:
-            (n x d) array of generated points.
+            (n x tunable_d) array of generated points.
 
         """
         raise NotImplementedError("Base RandomGenerator can't generate samples.")
@@ -359,7 +359,6 @@ def _convert_bounds(self, bounds: list[tuple[float, float]]) -> Tensor | None:
 
             Args:
                 bounds: A list of (lower, upper) tuples for each column of X.
-                    Defined on [0, 1]^d.
 
             Returns:
                 Optional 2 x d tensor representing the bounds

ax/generators/random/sobol.py

@@ -126,14 +126,15 @@ def gen(
             self.init_position = none_throws(self.engine).num_generated
         return points, weights
 
-    def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray:
+    def _gen_samples(self, n: int, tunable_d: int, bounds: npt.NDArray) -> npt.NDArray:
         """Generate n samples.
 
         Args:
             n: Number of samples to generate.
             tunable_d: The dimension of the generated samples. This must
                 match the tunable parameters used while initializing the
                 Sobol engine.
+            bounds: A (tunable_d x 2) array of (lower, upper) bounds.
 
         Returns:
             A numpy array of samples of shape `(n x tunable_d)`.
@@ -144,4 +145,8 @@ def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray:
             raise ValueError(
                 "Sobol Engine must be initialized before candidate generation."
             )
-        return none_throws(self.engine).draw(n, dtype=torch.double).numpy()
+        unit_samples = none_throws(self.engine).draw(n, dtype=torch.double).numpy()
+        # Rescale from [0, 1]^d to [lower, upper]^d
+        lower = bounds[:, 0]
+        upper = bounds[:, 1]
+        return unit_samples * (upper - lower) + lower

ax/generators/random/uniform.py

@@ -41,16 +41,21 @@ def __init__(
             # Fast-forward the random state by generating & discarding samples.
             self._rs.uniform(size=(self.init_position))
 
-    def _gen_samples(self, n: int, tunable_d: int) -> npt.NDArray:
+    def _gen_samples(self, n: int, tunable_d: int, bounds: npt.NDArray) -> npt.NDArray:
         """Generate samples from the scipy uniform distribution.
 
         Args:
             n: Number of samples to generate.
             tunable_d: Dimension of samples to generate.
+            bounds: A (tunable_d x 2) array of (lower, upper) bounds.
 
         Returns:
-            samples: An (n x d) array of random points.
+            samples: An (n x tunable_d) array of random points.
 
         """
         self.init_position += n * tunable_d
-        return self._rs.uniform(size=(n, tunable_d))
+        return self._rs.uniform(
+            low=bounds[:, 0],
+            high=bounds[:, 1],
+            size=(n, tunable_d),
+        )

ax/generators/tests/test_random.py

@@ -37,12 +37,14 @@ def test_state(self) -> None:
 
     def test_RandomGeneratorGenSamples(self) -> None:
         with self.assertRaises(NotImplementedError):
-            self.random_model._gen_samples(n=1, tunable_d=1)
+            self.random_model._gen_samples(
+                n=1, tunable_d=1, bounds=np.array([[0.0, 1.0]])
+            )
 
     def test_RandomGeneratorGenUnconstrained(self) -> None:
         with self.assertRaises(NotImplementedError):
             self.random_model._gen_unconstrained(
-                n=1, d=2, tunable_feature_indices=np.array([])
+                n=1, d=2, tunable_feature_indices=np.array([], dtype=int)
             )
 
     def test_ConvertEqualityConstraints(self) -> None:

ax/generators/tests/test_sobol.py

@@ -306,11 +306,21 @@ def test_SobolGeneratorOnlineRestart(self) -> None:
             np.allclose(generated_points_single_batch, generated_points_two_trials)
         )
 
-    def test_SobolGeneratorBadBounds(self) -> None:
-        generator = SobolGenerator()
-        ssd = SearchSpaceDigest(feature_names=["x"], bounds=[(-1, 1)])
-        with self.assertRaisesRegex(ValueError, "This generator operates on"):
-            generator.gen(n=1, search_space_digest=ssd, rounding_func=lambda x: x)
+    def test_SobolGeneratorGeneralBounds(self) -> None:
+        generator = SobolGenerator(seed=0)
+        bounds = [(2.0, 10.0), (0.5, 5.5), (-3.0, 3.0)]
+        ssd = SearchSpaceDigest(
+            feature_names=["x0", "x1", "x2"],
+            bounds=bounds,
+        )
+        generated_points, weights = generator.gen(
+            n=5, search_space_digest=ssd, rounding_func=lambda x: x
+        )
+        self.assertEqual(np.shape(generated_points), (5, 3))
+        np_bounds = np.array(bounds)
+        self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
+        self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
+        self.assertTrue(np.all(weights == 1.0))
 
     def test_SobolGeneratorMaxDraws(self) -> None:
         generator = SobolGenerator(seed=0)

ax/generators/tests/test_uniform.py

@@ -167,13 +167,18 @@ def test_with_linear_constraints(self) -> None:
         self.assertTrue(np.shape(expected_points) == np.shape(generated_points))
         self.assertTrue(np.allclose(expected_points, generated_points))
 
-    def test_with_bad_bounds(self) -> None:
-        generator = UniformGenerator()
-        with self.assertRaises(ValueError):
-            generated_points, weights = generator.gen(
-                n=1,
-                search_space_digest=SearchSpaceDigest(
-                    feature_names=["a"], bounds=[(-1, 1)]
-                ),
-                rounding_func=lambda x: x,
-            )
+    def test_with_general_bounds(self) -> None:
+        generator = UniformGenerator(seed=self.seed)
+        bounds = [(2.0, 10.0), (0.5, 5.5), (-3.0, 3.0)]
+        ssd = SearchSpaceDigest(
+            feature_names=["x0", "x1", "x2"],
+            bounds=bounds,
+        )
+        generated_points, weights = generator.gen(
+            n=5, search_space_digest=ssd, rounding_func=lambda x: x
+        )
+        self.assertEqual(np.shape(generated_points), (5, 3))
+        np_bounds = np.array(bounds)
+        self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
+        self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
+        self.assertTrue(np.all(weights == 1.0))

ax/generators/utils.py

@@ -70,8 +70,7 @@ def rejection_sample(
     rounding_func: Callable[[npt.NDArray], npt.NDArray] | None = None,
     existing_points: npt.NDArray | None = None,
 ) -> tuple[npt.NDArray, int]:
-    """Rejection sample in parameter space. Parameter space is typically
-    [0, 1] for all tunable parameters.
+    """Rejection sample in parameter space.
 
     Generators must implement a `gen_unconstrained` method in order to support
     rejection sampling via this utility.
@@ -273,30 +272,6 @@ def tunable_feature_indices(
     return np.delete(feature_indices, fixed_feature_indices)
 
 
-def validate_bounds(
-    bounds: Sequence[tuple[float, float]],
-    fixed_feature_indices: npt.NDArray,
-) -> None:
-    """Ensure the requested space is [0,1]^d.
-
-    Args:
-        bounds: A list of d (lower, upper) tuples for each column of X.
-        fixed_feature_indices: Indices of features which are fixed at a
-            particular value.
-    """
-    for feature_idx, bound in enumerate(bounds):
-        # Bounds for fixed features are not unit-transformed.
-        if feature_idx in fixed_feature_indices:
-            continue
-
-        if bound[0] != 0 or bound[1] != 1:
-            raise ValueError(
-                "This generator operates on [0,1]^d. Please make use "
-                "of the UnitX transform in the Adapter, and ensure "
-                "task features are fixed."
-            )
-
-
 def best_observed_point(
     model: TorchGeneratorLike,
     bounds: Sequence[tuple[float, float]],