Skip to content

Commit

Permalink
Use updated BoTorch HitAndRunPolytopeSampler for Sobol fallback (face…
Browse files Browse the repository at this point in the history 
…book#2492)

Summary:
Addresses the issues reported in facebook#2373 by using the updated `HitAndRunPolytopeSampler` from pytorch/botorch#2358 (this will require this to land in nightly botorch before tests can pass).


Differential Revision: D58070591

Pulled By: Balandat
  • Loading branch information
@Balandat @facebook-github-bot
Balandat authored and facebook-github-bot committed Jun 5, 2024
1 parent fb1270d commit 78735cd
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 54 deletions.
50 changes: 26 additions & 24 deletions ax/models/random/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,29 +140,31 @@ def gen(
except SearchSpaceExhausted as e:
if self.fallback_to_sample_polytope:
logger.info(
"Rejection sampling exceeded specified maximum draws."
"Rejection sampling exceeded specified maximum draws. "
"Falling back on HitAndRunPolytopeSampler instead of "
f"{self.__class__.__name__}."
)
# If rejection sampling fails, try polytope sampler.
num_generated = (
len(self.generated_points)
if self.generated_points is not None
else 0
)
polytope_sampler = HitAndRunPolytopeSampler(
inequality_constraints=self._convert_inequality_constraints(
linear_constraints
linear_constraints,
),
equality_constraints=self._convert_equality_constraints(
len(bounds), fixed_features
d=len(bounds),
fixed_features=fixed_features,
),
interior_point=self._get_last_point(),
bounds=self._convert_bounds(bounds),
interior_point=self._get_last_point(),
n_burnin=100,
n_thinning=20,
seed=self.seed + num_generated,
)
num_generated = (
len(self.generated_points)
if self.generated_points is not None
else 0
)
points = polytope_sampler.draw(
n=n, seed=self.seed + num_generated
).numpy()
points = polytope_sampler.draw(n=n).numpy()
# TODO: Should this round & deduplicate?
else:
raise e
Expand Down Expand Up @@ -241,8 +243,8 @@ def _convert_inequality_constraints(
if linear_constraints is None:
return None
else:
A = torch.tensor(linear_constraints[0], dtype=torch.double)
b = torch.tensor(linear_constraints[1], dtype=torch.double)
A = torch.as_tensor(linear_constraints[0], dtype=torch.double)
b = torch.as_tensor(linear_constraints[1], dtype=torch.double)
return A, b

def _convert_equality_constraints(
Expand All @@ -263,16 +265,16 @@ def _convert_equality_constraints(
"""
if fixed_features is None:
return None
else:
n = len(fixed_features)
fixed_indices = sorted(fixed_features.keys())
fixed_vals = torch.tensor(
[fixed_features[i] for i in fixed_indices], dtype=torch.double
)
constraint_matrix = torch.zeros((n, d), dtype=torch.double)
for index in range(0, len(fixed_vals)):
constraint_matrix[index, fixed_indices[index]] = 1.0
return constraint_matrix, fixed_vals

n = len(fixed_features)
fixed_indices = sorted(fixed_features.keys())
fixed_vals = torch.tensor(
[fixed_features[i] for i in fixed_indices], dtype=torch.double
)
constraint_matrix = torch.zeros((n, d), dtype=torch.double)
for index in range(0, len(fixed_vals)):
constraint_matrix[index, fixed_indices[index]] = 1.0
return constraint_matrix, fixed_vals

def _convert_bounds(self, bounds: List[Tuple[float, float]]) -> Optional[Tensor]:
"""Helper method to convert bounds list used by the rejectionsampler to the
Expand Down
48 changes: 18 additions & 30 deletions ax/models/tests/test_sobol.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,8 @@ def test_SobolGeneratorAllTunable(self) -> None:
)
self.assertEqual(np.shape(generated_points), (3, 3))
np_bounds = np.array(bounds)
self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
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))
self.assertEqual(generator._get_state().get("init_position"), 3)

Expand All @@ -51,8 +51,8 @@ def test_SobolGeneratorFixedSpace(self) -> None:
)
self.assertEqual(np.shape(generated_points), (3, 2))
np_bounds = np.array(bounds)
self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
# Should error out if deduplicating since there's only one feasible point.
generator = SobolGenerator(seed=0, deduplicate=True)
with self.assertRaisesRegex(SearchSpaceExhausted, "Rejection sampling"):
Expand Down Expand Up @@ -83,8 +83,8 @@ def test_SobolGeneratorNoScramble(self) -> None:
)
self.assertEqual(np.shape(generated_points), (3, 4))
np_bounds = np.array(bounds)
self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))

def test_SobolGeneratorOnline(self) -> None:
# Verify that the generator will return the expected arms if called
Expand All @@ -108,8 +108,8 @@ def test_SobolGeneratorOnline(self) -> None:
rounding_func=lambda x: x,
)
self.assertEqual(weights, [1])
self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
self.assertTrue(generated_points[..., -1] == 1)
self.assertTrue(np.array_equal(expected_points, generated_points.flatten()))

Expand All @@ -129,7 +129,7 @@ def test_SobolGeneratorWithOrderConstraints(self) -> None:
rounding_func=lambda x: x,
)
self.assertEqual(np.shape(generated_points), (3, 4))
self.assertTrue(np.alltrue(generated_points[..., -1] == 0.5))
self.assertTrue(np.all(generated_points[..., -1] == 0.5))
self.assertTrue(
np.array_equal(
np.sort(generated_points[..., :-1], axis=-1),
Expand All @@ -149,16 +149,13 @@ def test_SobolGeneratorWithLinearConstraints(self) -> None:
generated_points, weights = generator.gen(
n=3,
bounds=bounds,
linear_constraints=(
A,
b,
),
linear_constraints=(A, b),
fixed_features={fixed_param_index: 1},
rounding_func=lambda x: x,
)
self.assertTrue(np.shape(generated_points) == (3, 4))
self.assertTrue(np.alltrue(generated_points[..., -1] == 1))
self.assertTrue(np.alltrue(generated_points @ A.transpose() <= b))
self.assertTrue(np.all(generated_points[..., -1] == 1))
self.assertTrue(np.all(generated_points @ A.transpose() <= b))

def test_SobolGeneratorFallbackToPolytopeSampler(self) -> None:
# Ten parameters with sum less than 1. In this example, the rejection
Expand All @@ -175,10 +172,7 @@ def test_SobolGeneratorFallbackToPolytopeSampler(self) -> None:
generated_points, weights = generator.gen(
n=3,
bounds=bounds,
linear_constraints=(
A,
b,
),
linear_constraints=(A, b),
rounding_func=lambda x: x,
)
# First call uses the original seed since no candidates are generated.
Expand All @@ -187,7 +181,7 @@ def test_SobolGeneratorFallbackToPolytopeSampler(self) -> None:
"exceeded specified maximum draws" in mock_logger.call_args[0][0]
)
self.assertTrue(np.shape(generated_points) == (3, 10))
self.assertTrue(np.alltrue(generated_points @ A.transpose() <= b))
self.assertTrue(np.all(generated_points @ A.transpose() <= b))

with mock.patch(
"botorch.utils.sampling.sample_polytope",
Expand All @@ -196,10 +190,7 @@ def test_SobolGeneratorFallbackToPolytopeSampler(self) -> None:
generator.gen(
n=3,
bounds=bounds,
linear_constraints=(
A,
b,
),
linear_constraints=(A, b),
rounding_func=lambda x: x,
)
# Second call uses seed 3 since 3 candidates are already generated.
Expand All @@ -216,16 +207,13 @@ def test_SobolGeneratorFallbackToPolytopeSamplerWithFixedParam(self) -> None:
generated_points, weights = generator.gen(
n=3,
bounds=bounds,
linear_constraints=(
A,
b,
),
linear_constraints=(A, b),
fixed_features={10: 1},
rounding_func=lambda x: x,
)
self.assertTrue(np.shape(generated_points) == (3, 11))
self.assertTrue(np.alltrue(generated_points[..., -1] == 1))
self.assertTrue(np.alltrue(generated_points @ A.transpose() <= b))
self.assertTrue(np.all(generated_points[..., -1] == 1))
self.assertTrue(np.all(generated_points @ A.transpose() <= b))

def test_SobolGeneratorOnlineRestart(self) -> None:
# Ensure a single batch generation can also equivalently done by
Expand Down

0 comments on commit 78735cd

Please sign in to comment.