def setUp(self):
qNEI_input_constructor = get_acqf_input_constructor(
qNoisyExpectedImprovement)
self.mock_input_constructor = mock.MagicMock(
qNEI_input_constructor, side_effect=qNEI_input_constructor)
# Adding wrapping here to be able to count calls and inspect arguments.
_register_acqf_input_constructor(
acqf_cls=DummyACQFClass,
input_constructor=self.mock_input_constructor,
)
self.botorch_model_class = SingleTaskGP
self.surrogate = Surrogate(
botorch_model_class=self.botorch_model_class)
self.X = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]])
self.Y = torch.tensor([[3.0], [4.0]])
self.Yvar = torch.tensor([[0.0], [2.0]])
self.training_data = TrainingData.from_block_design(X=self.X,
Y=self.Y,
Yvar=self.Yvar)
self.fidelity_features = [2]
self.surrogate.construct(training_data=self.training_data,
fidelity_features=self.fidelity_features)
self.search_space_digest = SearchSpaceDigest(
feature_names=["a", "b", "c"],
bounds=[(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)],
target_fidelities={2: 1.0},
)
self.botorch_acqf_class = DummyACQFClass
self.objective_weights = torch.tensor([1.0])
self.objective_thresholds = None
self.pending_observations = [torch.tensor([[1.0, 3.0, 4.0]])]
self.outcome_constraints = (torch.tensor([[1.0]]), torch.tensor([[0.5]
]))
self.linear_constraints = None
self.fixed_features = {1: 2.0}
self.options = {"best_f": 0.0}
self.acquisition = Acquisition(
botorch_acqf_class=self.botorch_acqf_class,
surrogate=self.surrogate,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
objective_thresholds=self.objective_thresholds,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
options=self.options,
)
self.inequality_constraints = [(torch.tensor([0, 1]),
torch.tensor([-1.0, 1.0]), 1)]
self.rounding_func = lambda x: x
self.optimizer_options = {
Keys.NUM_RESTARTS: 40,
Keys.RAW_SAMPLES: 1024
}
def best_out_of_sample_point(
self,
bounds: List[Tuple[float, float]],
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
fidelity_features: Optional[List[int]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
options: Optional[TConfig] = None,
) -> Tuple[Tensor, Tensor]:
"""Finds the best predicted point and the corresponding value of the
appropriate best point acquisition function.
"""
if fixed_features:
# When have fixed features, need `FixedFeatureAcquisitionFunction`
# which has peculiar instantiation (wraps another acquisition fn.),
# so need to figure out how to handle.
# TODO (ref: https://fburl.com/diff/uneqb3n9)
raise NotImplementedError("Fixed features not yet supported.")
options = options or {}
acqf_class, acqf_options = pick_best_out_of_sample_point_acqf_class(
outcome_constraints=outcome_constraints,
seed_inner=checked_cast_optional(int, options.get(Keys.SEED_INNER, None)),
qmc=checked_cast(bool, options.get(Keys.QMC, True)),
)
# Avoiding circular import between `Surrogate` and `Acquisition`.
from ax.models.torch.botorch_modular.acquisition import Acquisition
acqf = Acquisition( # TODO: For multi-fidelity, might need diff. class.
surrogate=self,
botorch_acqf_class=acqf_class,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
target_fidelities=target_fidelities,
options=acqf_options,
)
candidates, acqf_values = acqf.optimize(
# pyre-ignore[6]: Exp. Tensor, got List[Tuple[float, float]].
# TODO: Fix typing of `bounds` in `TorchModel`-s.
bounds=bounds,
n=1,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints
),
fixed_features=fixed_features,
)
return candidates[0], acqf_values[0]
def setUp(self):
self.botorch_model_class = SingleTaskGP
self.surrogate = Surrogate(
botorch_model_class=self.botorch_model_class)
self.X = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]])
self.Y = torch.tensor([[3.0], [4.0]])
self.Yvar = torch.tensor([[0.0], [2.0]])
self.training_data = TrainingData(X=self.X, Y=self.Y, Yvar=self.Yvar)
self.fidelity_features = [2]
self.surrogate.construct(training_data=self.training_data,
fidelity_features=self.fidelity_features)
self.bounds = [(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)]
self.botorch_acqf_class = DummyACQFClass
self.objective_weights = torch.tensor([1.0])
self.objective_thresholds = None
self.pending_observations = [
torch.tensor([[1.0, 3.0, 4.0]]),
torch.tensor([[2.0, 6.0, 8.0]]),
]
self.outcome_constraints = (torch.tensor([[1.0]]), torch.tensor([[0.5]
]))
self.linear_constraints = None
self.fixed_features = {1: 2.0}
self.target_fidelities = {2: 1.0}
self.options = {"best_f": 0.0}
self.acquisition = Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
objective_thresholds=self.objective_thresholds,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
target_fidelities=self.target_fidelities,
options=self.options,
)
self.inequality_constraints = [(torch.tensor([0, 1]),
torch.tensor([-1.0, 1.0]), 1)]
self.rounding_func = lambda x: x
self.optimizer_options = {
Keys.NUM_RESTARTS: 40,
Keys.RAW_SAMPLES: 1024
}
def test_init_moo(
self,
mock_get_X,
):
moo_training_data = TrainingData(
Xs=[self.X] * 3,
Ys=[self.Y] * 3,
Yvars=[self.Yvar] * 3,
)
moo_objective_weights = torch.tensor(
[-1.0, -1.0, 0.0],
)
moo_objective_thresholds = torch.tensor(
[0.5, 1.5, float("nan")],
)
self.surrogate.construct(
training_data=moo_training_data,
)
mock_get_X.return_value = (self.pending_observations[0], self.X[:1])
outcome_constraints = (
torch.tensor(
[[1.0, 0.0, 0.0]],
),
torch.tensor(
[[10.0]],
),
)
acquisition = Acquisition(
surrogate=self.surrogate,
botorch_acqf_class=qNoisyExpectedHypervolumeImprovement,
search_space_digest=self.search_space_digest,
objective_weights=moo_objective_weights,
pending_observations=self.pending_observations,
outcome_constraints=outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
options=self.options,
objective_thresholds=moo_objective_thresholds,
)
self.assertTrue(
torch.equal(
moo_objective_thresholds[:2], acquisition.objective_thresholds[:2]
)
)
self.assertTrue(np.isnan(acquisition.objective_thresholds[2].item()))
# test inferred objective_thresholds
with ExitStack() as es:
preds = torch.tensor(
[
[11.0, 2.0],
[9.0, 3.0],
],
)
es.enter_context(
mock.patch.object(
self.surrogate.model,
"posterior",
return_value=MockPosterior(
mean=preds,
samples=preds,
),
)
)
acquisition = Acquisition(
surrogate=self.surrogate,
search_space_digest=self.search_space_digest,
objective_weights=moo_objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
options=self.options,
)
self.assertTrue(
torch.equal(
acquisition.objective_thresholds[:2], torch.tensor([9.9, 3.3])
)
)
self.assertTrue(np.isnan(acquisition.objective_thresholds[2].item()))
def test_init(
self,
mock_botorch_acqf_class,
mock_compute_model_deps,
mock_get_objective_and_transform,
mock_subset_model,
mock_get_X,
):
with self.assertRaisesRegex(TypeError, ".* missing .* 'botorch_acqf_class'"):
Acquisition(
surrogate=self.surrogate,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
)
botorch_objective = LinearMCObjective(weights=torch.tensor([1.0]))
mock_get_objective_and_transform.return_value = (botorch_objective, None)
mock_get_X.return_value = (self.pending_observations[0], self.X[:1])
acquisition = Acquisition(
surrogate=self.surrogate,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
options=self.options,
objective_thresholds=self.objective_thresholds,
)
# Check `_get_X_pending_and_observed` kwargs
mock_get_X.assert_called_with(
Xs=[self.training_data.X],
pending_observations=self.pending_observations,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
bounds=self.search_space_digest.bounds,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
)
# Call `subset_model` only when needed
mock_subset_model.assert_called_with(
model=acquisition.surrogate.model,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
objective_thresholds=self.objective_thresholds,
)
mock_subset_model.reset_mock()
mock_get_objective_and_transform.reset_mock()
self.mock_input_constructor.reset_mock()
mock_botorch_acqf_class.reset_mock()
self.options[Keys.SUBSET_MODEL] = False
acquisition = Acquisition(
surrogate=self.surrogate,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
options=self.options,
)
mock_subset_model.assert_not_called()
# Check `get_botorch_objective_and_transform` kwargs
mock_get_objective_and_transform.assert_called_once()
_, ckwargs = mock_get_objective_and_transform.call_args
self.assertIs(ckwargs["model"], self.acquisition.surrogate.model)
self.assertIs(ckwargs["objective_weights"], self.objective_weights)
self.assertIs(ckwargs["outcome_constraints"], self.outcome_constraints)
self.assertTrue(torch.equal(ckwargs["X_observed"], self.X[:1]))
# Check final `acqf` creation
model_deps = {Keys.CURRENT_VALUE: 1.2}
self.mock_input_constructor.assert_called_once()
mock_botorch_acqf_class.assert_called_once()
_, ckwargs = self.mock_input_constructor.call_args
self.assertIs(ckwargs["model"], self.acquisition.surrogate.model)
self.assertIs(ckwargs["objective"], botorch_objective)
self.assertTrue(torch.equal(ckwargs["X_pending"], self.pending_observations[0]))
for k, v in chain(self.options.items(), model_deps.items()):
self.assertEqual(ckwargs[k], v)
def test_init(
self,
mock_botorch_acqf_class,
mock_compute_data_deps,
mock_compute_model_deps,
mock_get_objective,
mock_subset_model,
mock_get_X,
):
self.acquisition.default_botorch_acqf_class = None
with self.assertRaisesRegex(
ValueError, ".*`botorch_acqf_class` argument must be specified."
):
Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
)
botorch_objective = LinearMCObjective(weights=torch.tensor([1.0]))
mock_get_objective.return_value = botorch_objective
acquisition = Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
target_fidelities=self.target_fidelities,
options=self.options,
)
# Check `_get_X_pending_and_observed` kwargs
mock_get_X.assert_called_with(
Xs=[self.training_data.X],
pending_observations=self.pending_observations,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
bounds=self.bounds,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
)
# Call `subset_model` only when needed
mock_subset_model.assert_called_with(
acquisition.surrogate.model,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
)
mock_subset_model.reset_mock()
self.options[Keys.SUBSET_MODEL] = False
acquisition = Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
target_fidelities=self.target_fidelities,
options=self.options,
)
mock_subset_model.assert_not_called()
# Check `get_botorch_objective` kwargs
mock_get_objective.assert_called_with(
model=self.acquisition.surrogate.model,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
X_observed=torch.tensor([3.0]),
use_scalarized_objective=False,
)
# Check `compute_model_dependencies` kwargs
mock_compute_model_deps.assert_called_with(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
target_fidelities=self.target_fidelities,
options=self.options,
)
# Check `compute_data_dependencies` kwargs
mock_compute_data_deps.assert_called_with(training_data=self.training_data)
# Check final `acqf` creation
model_deps = {Keys.CURRENT_VALUE: 1.2}
data_deps = {}
mock_botorch_acqf_class.assert_called_with(
model=self.acquisition.surrogate.model,
objective=botorch_objective,
X_pending=torch.tensor([2.0]),
X_baseline=torch.tensor([3.0]),
**self.options,
**model_deps,
**data_deps,
)
def test_init(
self,
mock_botorch_acqf_class,
mock_compute_model_deps,
mock_get_objective,
mock_subset_model,
mock_get_X,
):
self.acquisition.default_botorch_acqf_class = None
with self.assertRaisesRegex(
ValueError,
".*`botorch_acqf_class` argument must be specified."):
Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
)
botorch_objective = LinearMCObjective(weights=torch.tensor([1.0]))
mock_get_objective.return_value = botorch_objective
mock_get_X.return_value = (self.pending_observations[0], self.X[:1])
acquisition = Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
target_fidelities=self.target_fidelities,
options=self.options,
)
# Check `_get_X_pending_and_observed` kwargs
mock_get_X.assert_called_with(
Xs=[self.training_data.X],
pending_observations=self.pending_observations,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
bounds=self.bounds,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
)
# Call `subset_model` only when needed
mock_subset_model.assert_called_with(
acquisition.surrogate.model,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
objective_thresholds=self.objective_thresholds,
)
mock_subset_model.reset_mock()
mock_get_objective.reset_mock()
mock_botorch_acqf_class.reset_mock()
self.options[Keys.SUBSET_MODEL] = False
acquisition = Acquisition(
surrogate=self.surrogate,
bounds=self.bounds,
objective_weights=self.objective_weights,
botorch_acqf_class=self.botorch_acqf_class,
pending_observations=self.pending_observations,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
target_fidelities=self.target_fidelities,
options=self.options,
)
mock_subset_model.assert_not_called()
# Check `get_botorch_objective` kwargs
mock_get_objective.assert_called_once()
_, ckwargs = mock_get_objective.call_args
self.assertIs(ckwargs["model"], self.acquisition.surrogate.model)
self.assertIs(ckwargs["objective_weights"], self.objective_weights)
self.assertIs(ckwargs["outcome_constraints"], self.outcome_constraints)
self.assertTrue(torch.equal(ckwargs["X_observed"], self.X[:1]))
self.assertFalse(ckwargs["use_scalarized_objective"])
# Check final `acqf` creation
model_deps = {Keys.CURRENT_VALUE: 1.2}
mock_botorch_acqf_class.assert_called_once()
_, ckwargs = mock_botorch_acqf_class.call_args
self.assertIs(ckwargs["model"], self.acquisition.surrogate.model)
self.assertIs(ckwargs["objective"], botorch_objective)
self.assertTrue(
torch.equal(ckwargs["X_pending"], self.pending_observations[0]))
for k, v in chain(self.options.items(), model_deps.items()):
self.assertEqual(ckwargs[k], v)