def test_choose_model_class_task_features(self):
# Only a single task feature can be used.
with self.assertRaisesRegex(NotImplementedError,
"Only a single task feature"):
choose_model_class(
Yvars=self.Yvars,
search_space_digest=SearchSpaceDigest(feature_names=[],
bounds=[],
task_features=[1, 2]),
)
# With fidelity features and unknown variances, use SingleTaskMultiFidelityGP.
self.assertEqual(
MultiTaskGP,
choose_model_class(
Yvars=self.none_Yvars,
search_space_digest=SearchSpaceDigest(feature_names=[],
bounds=[],
task_features=[1]),
),
)
# With fidelity features and known variances, use FixedNoiseMultiFidelityGP.
self.assertEqual(
FixedNoiseMultiTaskGP,
choose_model_class(
Yvars=self.Yvars,
search_space_digest=SearchSpaceDigest(feature_names=[],
bounds=[],
task_features=[1]),
),
)
def test_choose_model_class(self):
# Mix of known and unknown variances.
with self.assertRaisesRegex(
ValueError,
"Variances should all be specified, or none should be."):
choose_model_class(
Yvars=[torch.tensor([[0.0], [np.nan]])],
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
),
)
# Without fidelity/task features but with Yvar specifications, use FixedNoiseGP.
self.assertEqual(
FixedNoiseGP,
choose_model_class(
Yvars=self.Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
),
),
)
# W/out fidelity/task features and w/out Yvar specifications, use SingleTaskGP.
self.assertEqual(
SingleTaskGP,
choose_model_class(
Yvars=[torch.tensor([[float("nan")], [float("nan")]])],
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
),
),
)
def test_fixed_rank_BotorchModel(self, dtype=torch.float, cuda=False):
Xs1, Ys1, Yvars1, bounds, _, fns, __package__ = get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(dtype=dtype,
cuda=cuda,
constant_noise=True)
model = BotorchModel(multitask_gp_ranks={"y": 2, "w": 1})
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=[0],
),
metric_names=["y", "w"],
)
_mock_fit_model.assert_called_once()
# Check ranks
model_list = model.model.models
self.assertEqual(model_list[0]._rank, 2)
self.assertEqual(model_list[1]._rank, 1)
def setUp(self):
self.X = torch.tensor([[1.0, 0.0], [1.0, 1.0], [1.0, 3.0], [2.0, 2.0],
[3.0, 1.0]])
self.Y = torch.tensor([
[1.0, 0.0, 0.0],
[1.0, 1.0, 1.0],
[1.0, 3.0, 3.0],
[2.0, 2.0, 4.0],
[3.0, 1.0, 3.0],
])
self.Yvar = torch.zeros(5, 3)
self.outcome_constraints = (
torch.tensor([[0.0, 0.0, 1.0]]),
torch.tensor([[3.5]]),
)
self.objective_thresholds = torch.tensor([0.5, 1.5])
self.objective_weights = torch.tensor([1.0, 1.0])
bounds = [(0.0, 4.0), (0.0, 4.0)]
self.model = MultiObjectiveBotorchModel(model_predictor=dummy_predict)
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
self.model.fit(
Xs=[self.X],
Ys=[self.Y],
Yvars=[self.Yvar],
search_space_digest=SearchSpaceDigest(
feature_names=["x1", "x2"],
bounds=bounds,
),
metric_names=["a", "b", "c"],
)
_mock_fit_model.assert_called_once()
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.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.acquisition_options = {Keys.NUM_FANTASIES: 64}
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.objective_weights = torch.tensor([1.0])
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.options = {
Keys.FIDELITY_WEIGHTS: {2: 1.0},
Keys.COST_INTERCEPT: 1.0,
Keys.NUM_TRACE_OBSERVATIONS: 0,
}
def testRFModel(self):
Xs = [np.random.rand(10, 2) for i in range(2)]
Ys = [np.random.rand(10, 1) for i in range(2)]
Yvars = [np.random.rand(10, 1) for i in range(2)]
m = RandomForest(num_trees=5)
m.fit(
Xs=Xs,
Ys=Ys,
Yvars=Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=["x1", "x2"],
bounds=[(0, 1)] * 2,
),
metric_names=["y"],
)
self.assertEqual(len(m.models), 2)
self.assertEqual(len(m.models[0].estimators_), 5)
f, cov = m.predict(np.random.rand(5, 2))
self.assertEqual(f.shape, (5, 2))
self.assertEqual(cov.shape, (5, 2, 2))
f, cov = m.cross_validate(
Xs_train=Xs, Ys_train=Ys, Yvars_train=Yvars, X_test=np.random.rand(3, 2)
)
self.assertEqual(f.shape, (3, 2))
self.assertEqual(cov.shape, (3, 2, 2))
def test_BotorchModelConstraints(self):
Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = get_torch_test_data(
dtype=torch.float, cuda=False, constant_noise=True)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(dtype=torch.float,
cuda=False,
constant_noise=True)
# make infeasible
Xs2[0] = -1 * Xs2[0]
objective_weights = torch.tensor([-1.0, 1.0],
dtype=torch.float,
device=torch.device("cpu"))
n = 3
model = BotorchModel()
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
_mock_fit_model.assert_called_once()
# because there are no feasible points:
with self.assertRaises(ValueError):
model.gen(n, bounds, objective_weights)
def test_optimize_mixed(self, mock_optimize_acqf_mixed):
tkwargs = {
"dtype": self.acquisition.dtype,
"device": self.acquisition.device,
}
ssd = SearchSpaceDigest(
feature_names=["a", "b"],
bounds=[(0, 1), (0, 2)],
categorical_features=[1],
discrete_choices={1: [0, 1, 2]},
)
self.acquisition.optimize(
n=3,
search_space_digest=ssd,
inequality_constraints=self.inequality_constraints,
fixed_features=None,
rounding_func=self.rounding_func,
optimizer_options=self.optimizer_options,
)
mock_optimize_acqf_mixed.assert_called_with(
acq_function=self.acquisition.acqf,
bounds=mock.ANY,
q=3,
fixed_features_list=[{1: 0}, {1: 1}, {1: 2}],
inequality_constraints=self.inequality_constraints,
post_processing_func=self.rounding_func,
**self.optimizer_options,
)
# can't use assert_called_with on bounds due to ambiguous bool comparison
expected_bounds = torch.tensor(ssd.bounds, **tkwargs).transpose(0, 1)
self.assertTrue(
torch.equal(
mock_optimize_acqf_mixed.call_args[1]["bounds"], expected_bounds
)
)
def testSearchSpaceDigest(self):
# test required fields
with self.assertRaises(TypeError):
SearchSpaceDigest(bounds=[])
with self.assertRaises(TypeError):
SearchSpaceDigest(feature_names=[])
# test instantiation
ssd = SearchSpaceDigest(**self.kwargs)
self.assertEqual(dataclasses.asdict(ssd), self.kwargs)
# test default instatiation
for arg in self.kwargs:
if arg in {"feature_names", "bounds"}:
continue
ssd = SearchSpaceDigest(
**{k: v
for k, v in self.kwargs.items() if k != arg})
def gen(
self,
n: int,
bounds: List[Tuple[float, float]],
objective_weights: Tensor,
objective_thresholds: Optional[Tensor] = None,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata,
Optional[List[TCandidateMetadata]]]:
if self._search_space_digest is None:
raise RuntimeError("Must `fit` the model before calling `gen`.")
acq_options, opt_options = construct_acquisition_and_optimizer_options(
acqf_options=self.acquisition_options,
model_gen_options=model_gen_options)
# update bounds / target fidelities
new_ssd_args = {
**dataclasses.asdict(self._search_space_digest),
"bounds": bounds,
"target_fidelities": target_fidelities or {},
}
search_space_digest = SearchSpaceDigest(**new_ssd_args)
acqf = self._instantiate_acquisition(
search_space_digest=search_space_digest,
objective_weights=objective_weights,
objective_thresholds=objective_thresholds,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
acq_options=acq_options,
)
botorch_rounding_func = get_rounding_func(rounding_func)
candidates, expected_acquisition_value = acqf.optimize(
n=n,
search_space_digest=search_space_digest,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
optimizer_options=checked_cast(dict, opt_options),
)
gen_metadata: TGenMetadata = {
Keys.EXPECTED_ACQF_VAL: expected_acquisition_value.tolist()
}
if objective_weights.nonzero().numel() > 1:
gen_metadata["objective_thresholds"] = acqf.objective_thresholds
gen_metadata["objective_weights"] = acqf.objective_weights
return (
candidates.detach().cpu(),
torch.ones(n, dtype=self.surrogate.dtype),
gen_metadata,
None,
)
def test_evaluate_acquisition_function(self, _mock_kg,
_mock_construct_options):
model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=KnowledgeGradient,
acquisition_options=self.acquisition_options,
)
model.surrogate.construct(
training_data=self.training_data,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
fidelity_features=self.search_space_digest.fidelity_features,
),
)
model.evaluate_acquisition_function(
X=self.X,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
acq_options=self.acquisition_options,
)
# `_mock_kg` is a mock of class, so to check the mock `evaluate` on
# instance of that class, we use `_mock_kg.return_value.evaluate`
_mock_kg.return_value.evaluate.assert_called()
def setUp(self):
self.botorch_model_class = SingleTaskGP
self.mll_class = ExactMarginalLogLikelihood
self.device = torch.device("cpu")
self.dtype = torch.float
self.Xs, self.Ys, self.Yvars, self.bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype)
self.training_data = TrainingData.from_block_design(X=self.Xs[0],
Y=self.Ys[0],
Yvar=self.Yvars[0])
self.surrogate_kwargs = self.botorch_model_class.construct_inputs(
self.training_data)
self.surrogate = Surrogate(
botorch_model_class=self.botorch_model_class,
mll_class=self.mll_class)
self.search_space_digest = SearchSpaceDigest(
feature_names=["x1", "x2"],
bounds=self.bounds,
target_fidelities={1: 1.0},
)
self.metric_names = ["y"]
self.fixed_features = {1: 2.0}
self.refit = True
self.objective_weights = torch.tensor([-1.0, 1.0],
dtype=self.dtype,
device=self.device)
self.outcome_constraints = (torch.tensor([[1.0]]), torch.tensor([[0.5]
]))
self.linear_constraints = (
torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
torch.tensor([[0.5], [1.0]]),
)
self.options = {}
def fit(
self,
Xs: List[Tensor],
Ys: List[Tensor],
Yvars: List[Tensor],
search_space_digest: SearchSpaceDigest,
metric_names: List[str],
candidate_metadata: Optional[List[List[TCandidateMetadata]]] = None,
) -> None:
assert len(search_space_digest.task_features) == 0
assert len(search_space_digest.fidelity_features) == 0
for b in search_space_digest.bounds:
# REMBO assumes the input space is [-1, 1]^D
assert b == (-1, 1)
self.num_outputs = len(Xs)
# For convenience for now, assume X for all outcomes the same
X_D = _get_single_X(Xs)
X_d = self.project_down(X_D)
# Fit model in low-d space (adjusted to [0, 1]^d)
super().fit(
Xs=[self.to_01(X_d)] * self.num_outputs,
Ys=Ys,
Yvars=Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=[f"x{i}" for i in range(self.A.shape[1])],
bounds=[(0.0, 1.0)] * len(self.bounds_d),
task_features=search_space_digest.task_features,
fidelity_features=search_space_digest.fidelity_features,
),
metric_names=metric_names,
candidate_metadata=candidate_metadata,
)
def test_choose_model_class_fidelity_features(self):
# Only a single fidelity feature can be used.
with self.assertRaisesRegex(NotImplementedError,
"Only a single fidelity feature"):
choose_model_class(
Yvars=self.Yvars,
search_space_digest=SearchSpaceDigest(feature_names=[],
bounds=[],
fidelity_features=[1,
2]),
)
# No support for non-empty task & fidelity features yet.
with self.assertRaisesRegex(NotImplementedError,
"Multi-task multi-fidelity"):
choose_model_class(
Yvars=self.Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
task_features=[1],
fidelity_features=[1],
),
)
# With fidelity features and unknown variances, use SingleTaskMultiFidelityGP.
self.assertEqual(
SingleTaskMultiFidelityGP,
choose_model_class(
Yvars=self.none_Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
fidelity_features=[2],
),
),
)
# With fidelity features and known variances, use FixedNoiseMultiFidelityGP.
self.assertEqual(
FixedNoiseMultiFidelityGP,
choose_model_class(
Yvars=self.Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
fidelity_features=[2],
),
),
)
def test_BotorchMOOModel_with_random_scalarization_and_outcome_constraints(
self, dtype=torch.float, cuda=False
):
tkwargs = {
"device": torch.device("cuda") if cuda else torch.device("cpu"),
"dtype": dtype,
}
Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True
)
Xs2, Ys2, Yvars2, _, _, _, _ = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True
)
n = 2
objective_weights = torch.tensor([1.0, 1.0], **tkwargs)
obj_t = torch.tensor([1.0, 1.0], **tkwargs)
model = MultiObjectiveBotorchModel(acqf_constructor=get_NEI)
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
_mock_fit_model.assert_called_once()
with mock.patch(
SAMPLE_SIMPLEX_UTIL_PATH,
autospec=True,
return_value=torch.tensor([0.7, 0.3], **tkwargs),
) as _mock_sample_simplex, mock.patch(
"ax.models.torch.botorch_moo_defaults.optimize_acqf_list",
return_value=(X_dummy, acqfv_dummy),
) as _:
model.gen(
n,
bounds,
objective_weights,
outcome_constraints=(
torch.tensor([[1.0, 1.0]], **tkwargs),
torch.tensor([[10.0]], **tkwargs),
),
model_gen_options={
"acquisition_function_kwargs": {"random_scalarization": True},
"optimizer_kwargs": _get_optimizer_kwargs(),
},
objective_thresholds=obj_t,
)
self.assertEqual(n, _mock_sample_simplex.call_count)
def test_BotorchMOOModel_with_chebyshev_scalarization_and_outcome_constraints(
self, dtype=torch.float, cuda=False
):
tkwargs = {
"device": torch.device("cuda") if cuda else torch.device("cpu"),
"dtype": torch.float,
}
Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True
)
Xs2, Ys2, Yvars2, _, _, _, _ = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True
)
n = 2
objective_weights = torch.tensor([1.0, 1.0], **tkwargs)
model = MultiObjectiveBotorchModel(acqf_constructor=get_NEI)
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
_mock_fit_model.assert_called_once()
with mock.patch(
CHEBYSHEV_SCALARIZATION_PATH, wraps=get_chebyshev_scalarization
) as _mock_chebyshev_scalarization, mock.patch(
"ax.models.torch.botorch_defaults.optimize_acqf",
return_value=(X_dummy, acqfv_dummy),
) as _:
model.gen(
n,
bounds,
objective_weights,
outcome_constraints=(
torch.tensor([[1.0, 1.0]], **tkwargs),
torch.tensor([[10.0]], **tkwargs),
),
model_gen_options={
"acquisition_function_kwargs": {"chebyshev_scalarization": True},
"optimizer_kwargs": _get_optimizer_kwargs(),
},
)
# get_chebyshev_scalarization should be called once for generated candidate.
self.assertEqual(n, _mock_chebyshev_scalarization.call_count)
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 test_fit(self, mock_fit_gpytorch, mock_MLL, mock_state_dict):
surrogate = ListSurrogate(
botorch_submodel_class_per_outcome=self.
botorch_submodel_class_per_outcome,
mll_class=SumMarginalLogLikelihood,
)
# Checking that model is None before `fit` (and `construct`) calls.
self.assertIsNone(surrogate._model)
# Should instantiate mll and `fit_gpytorch_model` when `state_dict`
# is `None`.
surrogate.fit(
training_data=self.training_data,
search_space_digest=SearchSpaceDigest(
feature_names=self.feature_names,
bounds=self.bounds,
task_features=self.task_features,
),
metric_names=self.outcomes,
)
mock_state_dict.assert_not_called()
mock_MLL.assert_called_once()
mock_fit_gpytorch.assert_called_once()
mock_state_dict.reset_mock()
mock_MLL.reset_mock()
mock_fit_gpytorch.reset_mock()
# Should `load_state_dict` when `state_dict` is not `None`
# and `refit` is `False`.
state_dict = {"state_attribute": "value"}
surrogate.fit(
training_data=self.training_data,
search_space_digest=SearchSpaceDigest(
feature_names=self.feature_names,
bounds=self.bounds,
task_features=self.task_features,
),
metric_names=self.outcomes,
refit=False,
state_dict=state_dict,
)
mock_state_dict.assert_called_once()
mock_MLL.assert_not_called()
mock_fit_gpytorch.assert_not_called()
def testMkDiscreteChoices(self):
ssd1 = SearchSpaceDigest(
feature_names=["a", "b"],
bounds=[(0, 1), (0, 2)],
ordinal_features=[1],
discrete_choices={1: [0, 1, 2]},
)
dc1 = mk_discrete_choices(ssd1)
self.assertEqual(dc1, {1: [0, 1, 2]})
dc1_ff = mk_discrete_choices(ssd1, fixed_features={1: 0})
self.assertEqual(dc1_ff, {1: [0]})
ssd2 = SearchSpaceDigest(
feature_names=["a", "b", "c"],
bounds=[(0, 1), (0, 2), (3, 4)],
ordinal_features=[1],
categorical_features=[2],
discrete_choices={1: [0, 1, 2], 2: [3, 4]},
)
dc2_ff = mk_discrete_choices(ssd2, fixed_features={1: 0})
self.assertEqual(dc2_ff, {1: [0], 2: [3, 4]})
def testTorchModelUpdate(self):
model = TorchModel()
with self.assertRaises(NotImplementedError):
model.update(
Xs=[np.array(0)],
Ys=[np.array(0)],
Yvars=[np.array(1)],
search_space_digest=SearchSpaceDigest(feature_names=[],
bounds=[]),
metric_names=[],
)
def test_BotorchMOOModel_with_ehvi_and_outcome_constraints(
self, dtype=torch.float, cuda=False
):
tkwargs = {
"device": torch.device("cuda") if cuda else torch.device("cpu"),
"dtype": dtype,
}
Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True
)
Xs2, Ys2, Yvars2, _, _, _, _ = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True
)
n = 3
objective_weights = torch.tensor([1.0, 1.0], **tkwargs)
model = MultiObjectiveBotorchModel(acqf_constructor=get_EHVI)
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
_mock_fit_model.assert_called_once()
with mock.patch(
EHVI_ACQF_PATH, wraps=moo_monte_carlo.qExpectedHypervolumeImprovement
) as _mock_ehvi_acqf, mock.patch(
"ax.models.torch.botorch_defaults.optimize_acqf",
return_value=(X_dummy, acqfv_dummy),
) as _:
model.gen(
n,
bounds,
objective_weights,
outcome_constraints=(
torch.tensor([[1.0, 1.0]], **tkwargs),
torch.tensor([[10.0]], **tkwargs),
),
model_gen_options={"optimizer_kwargs": _get_optimizer_kwargs()},
objective_thresholds=torch.tensor([1.0, 1.0]),
)
# the EHVI acquisition function should be created only once.
self.assertEqual(1, _mock_ehvi_acqf.call_count)
def test_evaluate_acquisition_function(self, _, mock_torch_model):
ma = TorchModelBridge(
experiment=None,
search_space=None,
data=None,
model=None,
transforms=[],
torch_dtype=torch.float64,
torch_device=torch.device("cpu"),
)
# These attributes would've been set by `ArrayModelBridge` __init__, but it's
# mocked.
ma.model = mock_torch_model()
t = mock.MagicMock(Transform, autospec=True)
t.transform_observation_features.return_value = [
ObservationFeatures(parameters={
"x": 3.0,
"y": 4.0
})
]
ma.transforms = {"ExampleTransform": t}
ma.parameters = ["x", "y"]
model_eval_acqf = mock_torch_model.return_value.evaluate_acquisition_function
model_eval_acqf.return_value = torch.tensor([5.0], dtype=torch.float64)
acqf_vals = ma.evaluate_acquisition_function(
observation_features=[
ObservationFeatures(parameters={
"x": 1.0,
"y": 2.0
})
],
search_space_digest=SearchSpaceDigest(feature_names=[], bounds=[]),
objective_weights=np.array([1.0]),
objective_thresholds=None,
outcome_constraints=None,
linear_constraints=None,
fixed_features=None,
pending_observations=None,
)
self.assertEqual(acqf_vals, [5.0])
t.transform_observation_features.assert_called_with(
[ObservationFeatures(parameters={
"x": 1.0,
"y": 2.0
})])
model_eval_acqf.assert_called_once()
self.assertTrue(
torch.equal( # `call_args` is an (args, kwargs) tuple
model_eval_acqf.call_args[1]["X"],
torch.tensor([[3.0, 4.0]], dtype=torch.float64),
))
def testNumpyModelFit(self):
numpy_model = NumpyModel()
numpy_model.fit(
Xs=[np.array(0)],
Ys=[np.array(0)],
Yvars=[np.array(1)],
search_space_digest=SearchSpaceDigest(
feature_names=["x"],
bounds=[(0, 1)],
),
metric_names=["y"],
)
def testTorchModelFit(self):
torch_model = TorchModel()
torch_model.fit(
Xs=[np.array(0)],
Ys=[np.array(0)],
Yvars=[np.array(1)],
search_space_digest=SearchSpaceDigest(
feature_names=["x1"],
bounds=[(0, 1)],
),
metric_names=["y"],
)
def testTorchModelCrossValidate(self):
torch_model = TorchModel()
with self.assertRaises(NotImplementedError):
torch_model.cross_validate(
Xs_train=[np.array([1])],
Ys_train=[np.array([1])],
Yvars_train=[np.array([1])],
X_test=np.array([1]),
search_space_digest=SearchSpaceDigest(feature_names=[],
bounds=[]),
metric_names=[],
)
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, 2.0], [4.0, 3.0, 1.0]])
self.Yvar = torch.tensor([[0.0, 2.0, 1.0], [2.0, 0.0, 1.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)
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, -1.0, 0.0])
self.objective_thresholds = torch.tensor([2.0, 1.0, float("nan")])
self.pending_observations = [
torch.tensor([[1.0, 3.0, 4.0]]),
torch.tensor([[1.0, 3.0, 4.0]]),
torch.tensor([[1.0, 3.0, 4.0]]),
]
self.outcome_constraints = (
torch.tensor([[1.0, 0.5, 0.5]]),
torch.tensor([[0.5]]),
)
self.con_tfs = get_outcome_constraint_transforms(
self.outcome_constraints)
self.linear_constraints = None
self.fixed_features = {1: 2.0}
self.options = {}
self.acquisition = MOOAcquisition(
surrogate=self.surrogate,
search_space_digest=self.search_space_digest,
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,
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 extract_search_space_digest(search_space: SearchSpace,
param_names: List[str]) -> SearchSpaceDigest:
"""Extract basic parameter prpoerties from a search space."""
bounds: List[Tuple[Union[int, float], Union[int, float]]] = []
ordinal_features: List[int] = []
categorical_features: List[int] = []
discrete_choices: Dict[int, List[Union[int, float]]] = {}
task_features: List[int] = []
fidelity_features: List[int] = []
target_fidelities: Dict[int, Union[int, float]] = {}
for i, p_name in enumerate(param_names):
p = search_space.parameters[p_name]
if isinstance(p, ChoiceParameter):
if p.is_task:
task_features.append(i)
elif p.is_ordered:
ordinal_features.append(i)
else:
categorical_features.append(i)
# at this point we can assume that values are numeric due to transforms
discrete_choices[i] = p.values # pyre-ignore [6]
bounds.append((min(p.values), max(p.values))) # pyre-ignore [6]
elif isinstance(p, RangeParameter):
if p.log_scale:
raise ValueError(f"{p} is log scale")
if p.parameter_type == ParameterType.INT:
ordinal_features.append(i)
d_choices = list(range(int(p.lower), int(p.upper) + 1))
discrete_choices[i] = d_choices # pyre-ignore [6]
bounds.append((p.lower, p.upper))
else:
raise ValueError(f"Unknown parameter type {type(p)}")
if p.is_fidelity:
if not isinstance(not_none(p.target_value), (int, float)):
raise NotImplementedError(
"Only numerical target values are supported.")
target_fidelities[i] = checked_cast_to_tuple((int, float),
p.target_value)
fidelity_features.append(i)
return SearchSpaceDigest(
feature_names=param_names,
bounds=bounds,
ordinal_features=ordinal_features,
categorical_features=categorical_features,
discrete_choices=discrete_choices,
task_features=task_features,
fidelity_features=fidelity_features,
target_fidelities=target_fidelities,
)
def setUp(self):
self.botorch_model_class = SingleTaskGP
self.surrogate = Surrogate(
botorch_model_class=self.botorch_model_class)
self.acquisition_class = KnowledgeGradient
self.botorch_acqf_class = qKnowledgeGradient
self.acquisition_options = {Keys.NUM_FANTASIES: 64}
self.model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=self.acquisition_class,
acquisition_options=self.acquisition_options,
)
self.dtype = torch.float
Xs1, Ys1, Yvars1, self.bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(dtype=self.dtype,
offset=1.0)
self.Xs = Xs1 + Xs2
self.Ys = Ys1 + Ys2
self.Yvars = Yvars1 + Yvars2
self.X = Xs1[0]
self.Y = Ys1[0]
self.Yvar = Yvars1[0]
self.X2 = Xs2[0]
self.training_data = TrainingData(X=self.X, Y=self.Y, Yvar=self.Yvar)
self.search_space_digest = SearchSpaceDigest(
feature_names=["x1", "x2", "x3"],
bounds=[(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)],
task_features=[],
fidelity_features=[2],
target_fidelities={1: 1.0},
)
self.metric_names = ["y"]
self.metric_names_for_list_surrogate = ["y1", "y2"]
self.candidate_metadata = []
self.optimizer_options = {
Keys.NUM_RESTARTS: 40,
Keys.RAW_SAMPLES: 1024
}
self.model_gen_options = {
Keys.OPTIMIZER_KWARGS: self.optimizer_options
}
self.objective_weights = torch.tensor([1.0])
self.objective_thresholds = None
self.outcome_constraints = None
self.linear_constraints = None
self.fixed_features = None
self.pending_observations = None
self.rounding_func = "func"
def test_choose_model_class_discrete_features(self):
# With discrete features, use MixedSingleTaskyGP.
self.assertEqual(
MixedSingleTaskGP,
choose_model_class(
Yvars=self.none_Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
task_features=[],
categorical_features=[1],
),
),
)
def extract_search_space_digest(search_space: SearchSpace,
param_names: List[str]) -> SearchSpaceDigest:
"""Extract basic parameter prpoerties from a search space."""
bounds: List[Tuple[Union[int, float], Union[int, float]]] = []
ordinal_features: List[int] = []
categorical_features: List[int] = []
task_features: List[int] = []
fidelity_features: List[int] = []
target_fidelities: Dict[int, Union[int, float]] = {}
for i, p_name in enumerate(param_names):
p = search_space.parameters[p_name]
if isinstance(p, ChoiceParameter):
if p.parameter_type != ParameterType.INT:
raise ValueError(f"Choice parameter {p} not of integer type")
if p.is_task:
task_features.append(i)
elif p.is_ordered:
ordinal_features.append(i)
else:
categorical_features.append(i)
bounds.append(
(int(p.values[0]), int(p.values[-1]))) # pyre-ignore [6]
elif isinstance(p, RangeParameter):
if p.log_scale:
raise ValueError(f"{p} is log scale")
if p.parameter_type == ParameterType.INT:
task_features.append(i)
# TODO: modify to ordinal_features.append(i) and adjust models
# that rely on the current behavior
bounds.append((p.lower, p.upper))
else:
raise ValueError(f"Unknown parameter type {type(p)}")
if p.is_fidelity:
if not isinstance(not_none(p.target_value), (int, float)):
raise NotImplementedError(
"Only numerical target values are supported.")
target_fidelities[i] = checked_cast_to_tuple((int, float),
p.target_value)
fidelity_features.append(i)
return SearchSpaceDigest(
feature_names=param_names,
bounds=bounds,
ordinal_features=ordinal_features,
categorical_features=categorical_features,
task_features=task_features,
fidelity_features=fidelity_features,
target_fidelities=target_fidelities,
)
