def test_transform_callback_int_log(self, *_):
exp = get_branin_experiment(with_batch=True)
parameters = [
RangeParameter(
name="x1",
parameter_type=ParameterType.INT,
lower=1,
upper=100,
log_scale=True,
),
RangeParameter(
name="x2",
parameter_type=ParameterType.INT,
lower=1,
upper=100,
log_scale=True,
),
]
gpei = TorchModelBridge(
experiment=exp,
data=exp.fetch_data(),
search_space=SearchSpace(parameters=parameters),
model=BotorchModel(),
transforms=[IntToFloat, Log],
torch_dtype=torch.double,
fit_out_of_design=True,
)
transformed = gpei._transform_callback([0.5, 1.5])
self.assertTrue(np.allclose(transformed, [0.47712, 1.50515]))
def test_transform_callback_int(self, *_):
exp = get_branin_experiment(with_batch=True)
data = get_branin_data(trial_indices=exp.trials)
parameters = [
RangeParameter(name="x1",
parameter_type=ParameterType.INT,
lower=1,
upper=10),
RangeParameter(name="x2",
parameter_type=ParameterType.INT,
lower=5,
upper=15),
]
gpei = TorchModelBridge(
experiment=exp,
data=data,
search_space=SearchSpace(parameters=parameters),
model=BotorchModel(),
transforms=[IntToFloat],
torch_dtype=torch.double,
fit_out_of_design=True,
)
transformed = gpei._transform_callback([5.4, 7.6])
self.assertTrue(np.allclose(transformed, [5, 8]))
np_mb = ArrayModelBridge(
experiment=exp,
data=exp.fetch_data(),
search_space=SearchSpace(parameters=parameters),
model=NumpyModel(),
transforms=[IntToFloat],
)
transformed = np_mb._transform_callback(np.array([5.4, 7.6]))
self.assertTrue(np.allclose(transformed, [5, 8]))
def test_transform_callback_log(self, *_):
parameters = [
RangeParameter(
name="x1",
parameter_type=ParameterType.FLOAT,
lower=1,
upper=3,
log_scale=True,
),
RangeParameter(
name="x2",
parameter_type=ParameterType.FLOAT,
lower=1,
upper=3,
log_scale=True,
),
]
search_space = SearchSpace(parameters=parameters)
exp = get_branin_experiment(with_batch=True, search_space=search_space)
gpei = TorchModelBridge(
experiment=exp,
data=exp.fetch_data(),
search_space=search_space,
model=BotorchModel(),
transforms=[Log],
torch_dtype=torch.double,
fit_out_of_design=True,
)
transformed = gpei._transform_callback([1.2, 2.5])
self.assertTrue(np.allclose(transformed, [1.2, 2.5]))
def get_MTGP(
experiment: MultiTypeExperiment,
data: Data,
search_space: Optional[SearchSpace] = None,
) -> TorchModelBridge:
"""Instantiates a Multi-task GP model that generates points with EI."""
trial_index_to_type = {
t.index: t.trial_type
for t in experiment.trials.values()
}
return TorchModelBridge(
experiment=experiment,
search_space=search_space or experiment.search_space,
data=data,
model=BotorchModel(),
transforms=MTGP_trans,
transform_configs={
"TrialAsTask": {
"trial_level_map": {
"trial_type": trial_index_to_type
}
},
"ConvertMetricNames": tconfig_from_mt_experiment(experiment),
},
torch_dtype=torch.double,
torch_device=DEFAULT_TORCH_DEVICE,
)
def get_REMBO(
experiment: Experiment,
data: Data,
A: torch.Tensor,
initial_X_d: torch.Tensor,
bounds_d: List[Tuple[float, float]],
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
**model_kwargs: Any,
) -> TorchModelBridge:
"""Instantiates a BotorchModel."""
if search_space is None:
search_space = experiment.search_space
if data.df.empty: # pragma: no cover
raise ValueError("REMBO model requires non-empty data.")
return TorchModelBridge(
experiment=experiment,
search_space=search_space,
data=data,
model=REMBO(A=A,
initial_X_d=initial_X_d,
bounds_d=bounds_d,
**model_kwargs),
transforms=[CenteredUnitX, StandardizeY],
torch_dtype=dtype,
torch_device=device,
)
def get_tensor_converter_model(experiment: Experiment, data: Data) -> TorchModelBridge:
"""
Constructs a minimal model for converting things to tensors.
Model fitting will instantiate all of the transforms but will not do any
expensive (i.e. GP) fitting beyond that. The model will raise an error if
it is used for predicting or generating.
Will work for any search space regardless of types of parameters.
Args:
experiment: Experiment.
data: Data for fitting the model.
Returns: A torch modelbridge with transforms set.
"""
# Transforms is the minimal set that will work for converting any search
# space to tensors.
return TorchModelBridge(
experiment=experiment,
search_space=experiment.search_space,
data=data,
model=TorchModel(),
transforms=[Derelativize, SearchSpaceToChoice, OrderedChoiceEncode, IntToFloat],
transform_configs={
"Derelativize": {"use_raw_status_quo": True},
"SearchSpaceToChoice": {"use_ordered": True},
},
fit_out_of_design=True,
)
def get_botorch(
experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None,
dtype: torch.dtype = torch.double,
device: torch.device = DEFAULT_TORCH_DEVICE,
transforms: List[Type[Transform]] = Cont_X_trans + Y_trans,
model_constructor: TModelConstructor = get_and_fit_model, # pyre-ignore[9]
model_predictor: TModelPredictor = predict_from_model,
acqf_constructor: TAcqfConstructor = get_NEI, # pyre-ignore[9]
acqf_optimizer: TOptimizer = scipy_optimizer, # pyre-ignore[9]
refit_on_cv: bool = False,
refit_on_update: bool = True,
) -> TorchModelBridge:
"""Instantiates a BotorchModel."""
if search_space is None:
search_space = experiment.search_space
if data.df.empty: # pragma: no cover
raise ValueError("BotorchModel requires non-empty data.")
return TorchModelBridge(
experiment=experiment,
search_space=search_space,
data=data,
model=BotorchModel(
model_constructor=model_constructor,
model_predictor=model_predictor,
acqf_constructor=acqf_constructor,
acqf_optimizer=acqf_optimizer,
),
transforms=transforms,
torch_dtype=dtype,
torch_device=device,
)
def get_MTGP(
experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None,
trial_index: Optional[int] = None,
) -> TorchModelBridge:
"""Instantiates a Multi-task Gaussian Process (MTGP) model that generates
points with EI.
If the input experiment is a MultiTypeExperiment then a
Multi-type Multi-task GP model will be instantiated.
Otherwise, the model will be a Single-type Multi-task GP.
"""
if isinstance(experiment, MultiTypeExperiment):
trial_index_to_type = {
t.index: t.trial_type for t in experiment.trials.values()
}
transforms = MT_MTGP_trans
transform_configs = {
"TrialAsTask": {"trial_level_map": {"trial_type": trial_index_to_type}},
"ConvertMetricNames": tconfig_from_mt_experiment(experiment),
}
else:
# Set transforms for a Single-type MTGP model.
transforms = ST_MTGP_trans
transform_configs = None
# Choose the status quo features for the experiment from the selected trial.
# If trial_index is None, we will look for a status quo from the last
# experiment trial to use as a status quo for the experiment.
if trial_index is None:
trial_index = len(experiment.trials) - 1
elif trial_index >= len(experiment.trials):
raise ValueError("trial_index is bigger than the number of experiment trials")
# pyre-fixme[16]: `ax.core.base_trial.BaseTrial` has no attribute `status_quo`.
status_quo = experiment.trials[trial_index].status_quo
if status_quo is None:
status_quo_features = None
else:
status_quo_features = ObservationFeatures(
parameters=status_quo.parameters, trial_index=trial_index
)
return TorchModelBridge(
experiment=experiment,
search_space=search_space or experiment.search_space,
data=data,
model=BotorchModel(),
transforms=transforms,
transform_configs=transform_configs,
torch_dtype=torch.double,
torch_device=DEFAULT_TORCH_DEVICE,
status_quo_features=status_quo_features,
)
def test_transform_callback_unitx(self, *_):
exp = get_branin_experiment(with_batch=True)
parameters = [
RangeParameter(name="x1",
parameter_type=ParameterType.FLOAT,
lower=0,
upper=10),
RangeParameter(name="x2",
parameter_type=ParameterType.FLOAT,
lower=0,
upper=100),
]
gpei = TorchModelBridge(
experiment=exp,
data=exp.fetch_data(),
search_space=SearchSpace(parameters=parameters),
model=BotorchModel(),
transforms=[UnitX],
)
transformed = gpei._transform_callback([0.75, 0.35])
self.assertTrue(np.allclose(transformed, [0.75, 0.35]))
def test_transform_callback_int(self, _):
exp = get_branin_experiment()
parameters = [
RangeParameter(name="x1",
parameter_type=ParameterType.INT,
lower=1,
upper=10),
RangeParameter(name="x2",
parameter_type=ParameterType.INT,
lower=5,
upper=15),
]
gpei = TorchModelBridge(
experiment=exp,
data=exp.fetch_data(),
search_space=SearchSpace(parameters=parameters),
model=BotorchModel(),
transforms=[IntToFloat],
torch_dtype=torch.double,
)
transformed = gpei._transform_callback([5.4, 7.6])
self.assertTrue(np.allclose(transformed, [5, 8]))
def get_MTGP(
experiment: Experiment,
data: Data,
is_multi_type: bool = True,
search_space: Optional[SearchSpace] = None,
) -> TorchModelBridge:
"""Instantiates a Multi-task GP model that generates points with EI.
Args:
is_multi_type: If is_multi_type is True then experiment should be a
MultiTypeExperiment and a Multi-type Multi-task GP model will be
instantiated.
Otherwise, the model will be a Single-type Multi-task GP.
"""
if is_multi_type and isinstance(experiment, MultiTypeExperiment):
trial_index_to_type = {
t.index: t.trial_type
for t in experiment.trials.values()
}
transforms = MT_MTGP_trans
transform_configs = {
"TrialAsTask": {
"trial_level_map": {
"trial_type": trial_index_to_type
}
},
"ConvertMetricNames": tconfig_from_mt_experiment(experiment),
}
elif is_multi_type:
raise ValueError(
"If is_multi_type is True, the input experiment type should be "
"MultiTypeExperiment.")
else:
transforms = ST_MTGP_trans
transform_configs = None
return TorchModelBridge(
experiment=experiment,
search_space=search_space or experiment.search_space,
data=data,
model=BotorchModel(),
transforms=transforms,
transform_configs=transform_configs,
torch_dtype=torch.double,
torch_device=DEFAULT_TORCH_DEVICE,
)
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 get_ALEBO(
experiment: Experiment,
search_space: SearchSpace,
data: Data,
B: torch.Tensor,
**model_kwargs: Any,
) -> TorchModelBridge:
if search_space is None:
search_space = experiment.search_space
return TorchModelBridge(
experiment=experiment,
search_space=search_space,
data=data,
model=ALEBO(B=B, **model_kwargs),
transforms=ALEBO_X_trans + [Derelativize, StandardizeY], # pyre-ignore
torch_dtype=B.dtype,
torch_device=B.device,
)
def get_ALEBO_kernel_ablation(
experiment: Experiment,
search_space: SearchSpace,
data: Data,
B: torch.Tensor,
**model_kwargs: Any,
) -> TorchModelBridge:
if search_space is None:
search_space = experiment.search_space
return TorchModelBridge(
experiment=experiment,
search_space=search_space,
data=data,
model=ALEBO_kernel_ablation(B=B, **model_kwargs),
transforms=[CenteredUnitX, StandardizeY],
torch_dtype=B.dtype,
torch_device=B.device,
)
def get_MTGP(experiment: Experiment,
data: Data,
search_space: Optional[SearchSpace] = None) -> TorchModelBridge:
"""Instantiates a Multi-task Gaussian Process (MTGP) model that generates
points with EI.
If the input experiment is a MultiTypeExperiment then a
Multi-type Multi-task GP model will be instantiated.
Otherwise, the model will be a Single-type Multi-task GP.
"""
if isinstance(experiment, MultiTypeExperiment):
trial_index_to_type = {
t.index: t.trial_type
for t in experiment.trials.values()
}
transforms = MT_MTGP_trans
transform_configs = {
"TrialAsTask": {
"trial_level_map": {
"trial_type": trial_index_to_type
}
},
"ConvertMetricNames": tconfig_from_mt_experiment(experiment),
}
else:
# Set transforms for a Single-type MTGP model.
transforms = ST_MTGP_trans
transform_configs = None
return TorchModelBridge(
experiment=experiment,
search_space=search_space or experiment.search_space,
data=data,
model=BotorchModel(),
transforms=transforms,
transform_configs=transform_configs,
torch_dtype=torch.double,
torch_device=DEFAULT_TORCH_DEVICE,
)
def test_status_quo_for_non_monolithic_data(self):
exp = get_branin_experiment_with_multi_objective(with_status_quo=True)
sobol_generator = get_sobol(
search_space=exp.search_space,
)
sobol_run = sobol_generator.gen(n=5)
exp.new_batch_trial(sobol_run).set_status_quo_and_optimize_power(
status_quo=exp.status_quo
).run()
# create data where metrics vary in start and end times
data = get_non_monolithic_branin_moo_data()
bridge = TorchModelBridge(
search_space=exp.search_space,
model=MultiObjectiveBotorchModel(),
optimization_config=exp.optimization_config,
experiment=exp,
data=data,
transforms=[],
)
self.assertEqual(bridge.status_quo.arm_name, "status_quo")
def test_pareto_frontier(self, _):
exp = get_branin_experiment_with_multi_objective(
has_optimization_config=True, with_batch=True
)
for trial in exp.trials.values():
trial.mark_running(no_runner_required=True).mark_completed()
metrics_dict = exp.optimization_config.metrics
objective_thresholds = [
ObjectiveThreshold(
metric=metrics_dict["branin_a"],
bound=0.0,
relative=False,
op=ComparisonOp.GEQ,
),
ObjectiveThreshold(
metric=metrics_dict["branin_b"],
bound=0.0,
relative=False,
op=ComparisonOp.GEQ,
),
]
exp.optimization_config = exp.optimization_config.clone_with_args(
objective_thresholds=objective_thresholds
)
exp.attach_data(
get_branin_data_multi_objective(trial_indices=exp.trials.keys())
)
modelbridge = TorchModelBridge(
search_space=exp.search_space,
model=MultiObjectiveBotorchModel(),
optimization_config=exp.optimization_config,
transforms=[t1, t2],
experiment=exp,
data=exp.fetch_data(),
objective_thresholds=objective_thresholds,
)
with patch(
PARETO_FRONTIER_EVALUATOR_PATH, wraps=pareto_frontier_evaluator
) as wrapped_frontier_evaluator:
modelbridge.model.frontier_evaluator = wrapped_frontier_evaluator
observed_frontier = observed_pareto_frontier(
modelbridge=modelbridge, objective_thresholds=objective_thresholds
)
wrapped_frontier_evaluator.assert_called_once()
self.assertIsNone(wrapped_frontier_evaluator.call_args[1]["X"])
self.assertEqual(1, len(observed_frontier))
self.assertEqual(observed_frontier[0].arm_name, "0_0")
with self.assertRaises(ValueError):
predicted_pareto_frontier(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=[],
)
predicted_frontier = predicted_pareto_frontier(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=None,
)
self.assertEqual(predicted_frontier[0].arm_name, "0_0")
observation_features = [
ObservationFeatures(parameters={"x1": 0.0, "x2": 1.0}),
ObservationFeatures(parameters={"x1": 1.0, "x2": 0.0}),
]
observation_data = [
ObservationData(
metric_names=["branin_b", "branin_a"],
means=np.array([1.0, 2.0]),
covariance=np.array([[1.0, 2.0], [3.0, 4.0]]),
),
ObservationData(
metric_names=["branin_a", "branin_b"],
means=np.array([3.0, 4.0]),
covariance=np.array([[1.0, 2.0], [3.0, 4.0]]),
),
]
predicted_frontier = predicted_pareto_frontier(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=observation_features,
)
self.assertTrue(len(predicted_frontier) <= 2)
self.assertIsNone(predicted_frontier[0].arm_name, None)
with patch(
PARETO_FRONTIER_EVALUATOR_PATH, wraps=pareto_frontier_evaluator
) as wrapped_frontier_evaluator:
(observed_frontier, f, obj_w, obj_t,) = get_pareto_frontier_and_configs(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=observation_features,
observation_data=observation_data,
)
wrapped_frontier_evaluator.assert_called_once()
self.assertTrue(
torch.equal(
wrapped_frontier_evaluator.call_args[1]["X"],
torch.tensor([[1.0, 4.0], [4.0, 1.0]]),
)
)
self.assertEqual(f.shape, (1, 2))
self.assertTrue(torch.equal(obj_w, torch.tensor([1.0, 1.0])))
self.assertTrue(torch.equal(obj_t, torch.tensor([0.0, 0.0])))
observed_frontier2 = pareto_frontier(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=observation_features,
observation_data=observation_data,
)
self.assertEqual(observed_frontier, observed_frontier2)
with patch(
PARETO_FRONTIER_EVALUATOR_PATH, wraps=pareto_frontier_evaluator
) as wrapped_frontier_evaluator:
(observed_frontier, f, obj_w, obj_t,) = get_pareto_frontier_and_configs(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=observation_features,
observation_data=observation_data,
use_model_predictions=False,
)
wrapped_frontier_evaluator.assert_called_once()
self.assertIsNone(wrapped_frontier_evaluator.call_args[1]["X"])
true_Y = torch.tensor([[9.0, 4.0], [16.0, 25.0]])
self.assertTrue(
torch.equal(
wrapped_frontier_evaluator.call_args[1]["Y"],
true_Y,
)
)
self.assertTrue(torch.equal(f, true_Y[1:, :]))
def test_hypervolume(self, _, cuda=False):
for num_objectives in (2, 3):
exp = get_branin_experiment_with_multi_objective(
has_optimization_config=True,
with_batch=True,
num_objectives=num_objectives,
)
for trial in exp.trials.values():
trial.mark_running(no_runner_required=True).mark_completed()
metrics_dict = exp.optimization_config.metrics
objective_thresholds = [
ObjectiveThreshold(
metric=metrics_dict["branin_a"],
bound=0.0,
relative=False,
op=ComparisonOp.GEQ,
),
ObjectiveThreshold(
metric=metrics_dict["branin_b"],
bound=1.0,
relative=False,
op=ComparisonOp.GEQ,
),
]
if num_objectives == 3:
objective_thresholds.append(
ObjectiveThreshold(
metric=metrics_dict["branin_c"],
bound=2.0,
relative=False,
op=ComparisonOp.GEQ,
)
)
optimization_config = exp.optimization_config.clone_with_args(
objective_thresholds=objective_thresholds
)
exp.attach_data(
get_branin_data_multi_objective(
trial_indices=exp.trials.keys(), num_objectives=num_objectives
)
)
modelbridge = TorchModelBridge(
search_space=exp.search_space,
model=MultiObjectiveBotorchModel(),
optimization_config=optimization_config,
transforms=[],
experiment=exp,
data=exp.fetch_data(),
torch_device=torch.device("cuda" if cuda else "cpu"),
objective_thresholds=objective_thresholds,
)
with patch(
PARETO_FRONTIER_EVALUATOR_PATH, wraps=pareto_frontier_evaluator
) as wrapped_frontier_evaluator:
modelbridge.model.frontier_evaluator = wrapped_frontier_evaluator
hv = observed_hypervolume(
modelbridge=modelbridge, objective_thresholds=objective_thresholds
)
expected_hv = 20 if num_objectives == 2 else 60 # 5 * 4 (* 3)
wrapped_frontier_evaluator.assert_called_once()
self.assertEqual(expected_hv, hv)
if num_objectives == 3:
# Test selected_metrics
hv = observed_hypervolume(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
selected_metrics=["branin_a", "branin_c"],
)
expected_hv = 15 # (5 - 0) * (5 - 2)
self.assertEqual(expected_hv, hv)
# test that non-objective outcome raises value error
with self.assertRaises(ValueError):
hv = observed_hypervolume(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
selected_metrics=["tracking"],
)
with self.assertRaises(ValueError):
predicted_hypervolume(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=[],
)
observation_features = [
ObservationFeatures(parameters={"x1": 1.0, "x2": 2.0}),
ObservationFeatures(parameters={"x1": 2.0, "x2": 1.0}),
]
predicted_hv = predicted_hypervolume(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=observation_features,
)
self.assertTrue(predicted_hv >= 0)
if num_objectives == 3:
# Test selected_metrics
predicted_hv = predicted_hypervolume(
modelbridge=modelbridge,
objective_thresholds=objective_thresholds,
observation_features=observation_features,
selected_metrics=["branin_a", "branin_c"],
)
self.assertTrue(predicted_hv >= 0)
def test_infer_objective_thresholds(self, _, cuda=False):
# lightweight test
exp = get_branin_experiment_with_multi_objective(
has_optimization_config=True,
with_batch=True,
with_status_quo=True,
)
for trial in exp.trials.values():
trial.mark_running(no_runner_required=True).mark_completed()
exp.attach_data(
get_branin_data_multi_objective(trial_indices=exp.trials.keys())
)
data = exp.fetch_data()
modelbridge = TorchModelBridge(
search_space=exp.search_space,
model=MultiObjectiveBotorchModel(),
optimization_config=exp.optimization_config,
transforms=Cont_X_trans + Y_trans,
torch_device=torch.device("cuda" if cuda else "cpu"),
experiment=exp,
data=data,
)
fixed_features = ObservationFeatures(parameters={"x1": 0.0})
search_space = exp.search_space.clone()
param_constraints = [
ParameterConstraint(constraint_dict={"x1": 1.0}, bound=10.0)
]
search_space.add_parameter_constraints(param_constraints)
oc = exp.optimization_config.clone()
oc.objective._objectives[0].minimize = True
expected_base_gen_args = modelbridge._get_transformed_gen_args(
search_space=search_space.clone(),
optimization_config=oc,
fixed_features=fixed_features,
)
with ExitStack() as es:
mock_model_infer_obj_t = es.enter_context(
patch(
"ax.modelbridge.torch.infer_objective_thresholds",
wraps=infer_objective_thresholds,
)
)
mock_get_transformed_gen_args = es.enter_context(
patch.object(
modelbridge,
"_get_transformed_gen_args",
wraps=modelbridge._get_transformed_gen_args,
)
)
mock_get_transformed_model_gen_args = es.enter_context(
patch.object(
modelbridge,
"_get_transformed_model_gen_args",
wraps=modelbridge._get_transformed_model_gen_args,
)
)
mock_untransform_objective_thresholds = es.enter_context(
patch.object(
modelbridge,
"_untransform_objective_thresholds",
wraps=modelbridge._untransform_objective_thresholds,
)
)
obj_thresholds = modelbridge.infer_objective_thresholds(
search_space=search_space,
optimization_config=oc,
fixed_features=fixed_features,
)
expected_obj_weights = torch.tensor([-1.0, 1.0])
ckwargs = mock_model_infer_obj_t.call_args[1]
self.assertTrue(
torch.equal(ckwargs["objective_weights"], expected_obj_weights)
)
# check that transforms have been applied (at least UnitX)
self.assertEqual(ckwargs["bounds"], [(0.0, 1.0), (0.0, 1.0)])
lc = ckwargs["linear_constraints"]
self.assertTrue(torch.equal(lc[0], torch.tensor([[15.0, 0.0]])))
self.assertTrue(torch.equal(lc[1], torch.tensor([[15.0]])))
self.assertEqual(ckwargs["fixed_features"], {0: 1.0 / 3.0})
mock_get_transformed_gen_args.assert_called_once()
mock_get_transformed_model_gen_args.assert_called_once_with(
search_space=expected_base_gen_args.search_space,
fixed_features=expected_base_gen_args.fixed_features,
pending_observations=expected_base_gen_args.pending_observations,
optimization_config=expected_base_gen_args.optimization_config,
)
mock_untransform_objective_thresholds.assert_called_once()
ckwargs = mock_untransform_objective_thresholds.call_args[1]
self.assertTrue(
torch.equal(ckwargs["objective_weights"], expected_obj_weights)
)
self.assertEqual(ckwargs["bounds"], [(0.0, 1.0), (0.0, 1.0)])
self.assertEqual(ckwargs["fixed_features"], {0: 1.0 / 3.0})
self.assertEqual(obj_thresholds[0].metric.name, "branin_a")
self.assertEqual(obj_thresholds[1].metric.name, "branin_b")
self.assertEqual(obj_thresholds[0].op, ComparisonOp.LEQ)
self.assertEqual(obj_thresholds[1].op, ComparisonOp.GEQ)
self.assertFalse(obj_thresholds[0].relative)
self.assertFalse(obj_thresholds[1].relative)
df = exp_to_df(exp)
Y = np.stack([df.branin_a.values, df.branin_b.values]).T
Y = torch.from_numpy(Y)
Y[:, 0] *= -1
pareto_Y = Y[is_non_dominated(Y)]
nadir = pareto_Y.min(dim=0).values
self.assertTrue(
np.all(
np.array([-obj_thresholds[0].bound, obj_thresholds[1].bound])
< nadir.numpy()
)
)
# test using MTGP
sobol_generator = get_sobol(
search_space=exp.search_space,
seed=TEST_SOBOL_SEED,
# set initial position equal to the number of sobol arms generated
# so far. This means that new sobol arms will complement the previous
# arms in a space-filling fashion
init_position=len(exp.arms_by_name) - 1,
)
sobol_run = sobol_generator.gen(n=2)
trial = exp.new_batch_trial(optimize_for_power=True)
trial.add_generator_run(sobol_run)
trial.mark_running(no_runner_required=True).mark_completed()
data = exp.fetch_data()
torch.manual_seed(0) # make model fitting deterministic
modelbridge = TorchModelBridge(
search_space=exp.search_space,
model=MultiObjectiveBotorchModel(),
optimization_config=exp.optimization_config,
transforms=ST_MTGP_trans,
experiment=exp,
data=data,
)
fixed_features = ObservationFeatures(parameters={}, trial_index=1)
expected_base_gen_args = modelbridge._get_transformed_gen_args(
search_space=search_space.clone(),
optimization_config=exp.optimization_config,
fixed_features=fixed_features,
)
with ExitStack() as es:
mock_model_infer_obj_t = es.enter_context(
patch(
"ax.modelbridge.torch.infer_objective_thresholds",
wraps=infer_objective_thresholds,
)
)
mock_untransform_objective_thresholds = es.enter_context(
patch.object(
modelbridge,
"_untransform_objective_thresholds",
wraps=modelbridge._untransform_objective_thresholds,
)
)
obj_thresholds = modelbridge.infer_objective_thresholds(
search_space=search_space,
optimization_config=exp.optimization_config,
fixed_features=fixed_features,
)
ckwargs = mock_model_infer_obj_t.call_args[1]
self.assertEqual(ckwargs["fixed_features"], {2: 1.0})
mock_untransform_objective_thresholds.assert_called_once()
ckwargs = mock_untransform_objective_thresholds.call_args[1]
self.assertEqual(ckwargs["fixed_features"], {2: 1.0})
self.assertEqual(obj_thresholds[0].metric.name, "branin_a")
self.assertEqual(obj_thresholds[1].metric.name, "branin_b")
self.assertEqual(obj_thresholds[0].op, ComparisonOp.GEQ)
self.assertEqual(obj_thresholds[1].op, ComparisonOp.GEQ)
self.assertFalse(obj_thresholds[0].relative)
self.assertFalse(obj_thresholds[1].relative)
df = exp_to_df(exp)
trial_mask = df.trial_index == 1
Y = np.stack([df.branin_a.values[trial_mask], df.branin_b.values[trial_mask]]).T
Y = torch.from_numpy(Y)
pareto_Y = Y[is_non_dominated(Y)]
nadir = pareto_Y.min(dim=0).values
self.assertTrue(
np.all(
np.array([obj_thresholds[0].bound, obj_thresholds[1].bound])
< nadir.numpy()
)
)
def testTorchModelBridge(self, mock_init):
torch_dtype = torch.float64
torch_device = torch.device("cpu")
ma = TorchModelBridge(
experiment=None,
search_space=None,
data=None,
model=None,
transforms=[],
torch_dtype=torch.float64,
torch_device=torch.device("cpu"),
)
self.assertEqual(ma.dtype, torch.float64)
self.assertEqual(ma.device, torch.device("cpu"))
self.assertFalse(mock_init.call_args[-1]["fit_out_of_design"])
# Test `fit`.
model = mock.MagicMock(TorchModel, autospec=True, instance=True)
X = np.array([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]])
Y = np.array([[3.0], [4.0]])
var = np.array([[1.0], [2.0]])
ma._model_fit(
model=model,
Xs=[X],
Ys=[Y],
Yvars=[var],
search_space_digest=SearchSpaceDigest(feature_names=[], bounds=[]),
metric_names=[],
candidate_metadata=[],
)
model_fit_args = model.fit.mock_calls[0][2]
self.assertTrue(
torch.equal(
model_fit_args["Xs"][0],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_fit_args["Ys"][0],
torch.tensor(Y, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_fit_args["Yvars"][0],
torch.tensor(var, dtype=torch_dtype, device=torch_device),
))
# Test `update` (need to fill required fields before call to `_model_update`).
ma.parameters = []
ma.outcomes = []
ma._model_update(
Xs=[X],
Ys=[Y],
Yvars=[var],
search_space_digest=SearchSpaceDigest(feature_names=[], bounds=[]),
metric_names=[],
candidate_metadata=[],
)
model_update_args = model.update.mock_calls[0][2]
self.assertTrue(
torch.equal(
model_update_args["Xs"][0],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_update_args["Ys"][0],
torch.tensor(Y, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_update_args["Yvars"][0],
torch.tensor(var, dtype=torch_dtype, device=torch_device),
))
# Predict
model.predict.return_value = (torch.tensor([3.0]), torch.tensor([4.0]))
f, var = ma._model_predict(X)
self.assertTrue(
torch.equal(
model.predict.mock_calls[0][2]["X"],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(np.array_equal(f, np.array([3.0])))
self.assertTrue(np.array_equal(var, np.array([4.0])))
# Gen
model.gen.return_value = (
torch.tensor([1.0, 2.0, 3.0]),
torch.tensor([1.0]),
{},
[],
)
X, w, _gen_metadata, _candidate_metadata = ma._model_gen(
n=3,
bounds=[(0, 1)],
objective_weights=np.array([1.0, 0.0]),
outcome_constraints=None,
linear_constraints=None,
fixed_features={1: 3.0},
pending_observations=[np.array([]),
np.array([1.0, 2.0, 3.0])],
model_gen_options={"option": "yes"},
rounding_func=np.round,
target_fidelities=None,
)
gen_args = model.gen.mock_calls[0][2]
self.assertEqual(gen_args["n"], 3)
self.assertEqual(gen_args["bounds"], [(0, 1)])
self.assertTrue(
torch.equal(
gen_args["objective_weights"],
torch.tensor([1.0, 0.0],
dtype=torch_dtype,
device=torch_device),
))
self.assertIsNone(gen_args["outcome_constraints"])
self.assertIsNone(gen_args["linear_constraints"])
self.assertEqual(gen_args["fixed_features"], {1: 3.0})
self.assertTrue(
torch.equal(
gen_args["pending_observations"][0],
torch.tensor([], dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
gen_args["pending_observations"][1],
torch.tensor([1.0, 2.0, 3.0],
dtype=torch_dtype,
device=torch_device),
))
self.assertEqual(gen_args["model_gen_options"], {"option": "yes"})
self.assertIsNone(gen_args["target_fidelities"])
# check rounding function
t = torch.tensor([0.1, 0.6], dtype=torch_dtype, device=torch_device)
self.assertTrue(
torch.equal(gen_args["rounding_func"](t), torch.round(t)))
self.assertTrue(np.array_equal(X, np.array([1.0, 2.0, 3.0])))
self.assertTrue(np.array_equal(w, np.array([1.0])))
# Cross-validate
model.cross_validate.return_value = (torch.tensor([3.0]),
torch.tensor([4.0]))
f, var = ma._model_cross_validate(
Xs_train=[X],
Ys_train=[Y],
Yvars_train=[var],
X_test=X,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[(0, 1)],
),
metric_names=[],
)
model_cv_args = model.cross_validate.mock_calls[0][2]
self.assertTrue(
torch.equal(
model_cv_args["Xs_train"][0],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_cv_args["Ys_train"][0],
torch.tensor(Y, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_cv_args["Yvars_train"][0],
torch.tensor(var, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_cv_args["X_test"],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(np.array_equal(f, np.array([3.0])))
self.assertTrue(np.array_equal(var, np.array([4.0])))
# Transform observation features
obsf = [ObservationFeatures(parameters={"x": 1.0, "y": 2.0})]
ma.parameters = ["x", "y"]
X = ma._transform_observation_features(obsf)
self.assertTrue(
torch.equal(
X,
torch.tensor([[1.0, 2.0]],
dtype=torch_dtype,
device=torch_device)))
# test fit out of design
ma = TorchModelBridge(
experiment=None,
search_space=None,
data=None,
model=None,
transforms=[],
torch_dtype=torch.float64,
torch_device=torch.device("cpu"),
fit_out_of_design=True,
)
self.assertTrue(mock_init.call_args[-1]["fit_out_of_design"])
def get_observed_pareto_frontiers(
experiment: Experiment,
data: Optional[Data] = None,
rel: bool = True,
) -> List[ParetoFrontierResults]:
"""
Find all Pareto points from an experiment.
Uses only values as observed in the data; no modeling is involved. Makes no
assumption about the search space or types of parameters. If "data" is provided will
use that, otherwise will use all data attached to the experiment.
Uses all arms present in data; does not filter according to experiment
search space.
Assumes experiment has a multiobjective optimization config from which the
objectives and outcome constraints will be extracted.
Will generate a ParetoFrontierResults for every pair of metrics in the experiment's
multiobjective optimization config.
"""
if data is None:
data = experiment.fetch_data()
if experiment.optimization_config is None:
raise ValueError("Experiment must have an optimization config")
# Make a dummy model for converting things to tensors.
# Transforms is the minimal set that will work for converting any search
# space to tensors.
mb = TorchModelBridge(
experiment=experiment,
search_space=experiment.search_space,
data=data,
model=TorchModel(),
transforms=[Derelativize, SearchSpaceToChoice, OneHot],
transform_configs={"Derelativize": {
"use_raw_status_quo": True
}},
fit_out_of_design=True,
)
pareto_observations = observed_pareto_frontier(modelbridge=mb)
# Convert to ParetoFrontierResults
metric_names = [
metric.name for metric in
experiment.optimization_config.objective.metrics # pyre-ignore
]
pfr_means = {name: [] for name in metric_names}
pfr_sems = {name: [] for name in metric_names}
for obs in pareto_observations:
for i, name in enumerate(obs.data.metric_names):
pfr_means[name].append(obs.data.means[i])
pfr_sems[name].append(np.sqrt(obs.data.covariance[i, i]))
# Relativize as needed
if rel and experiment.status_quo is not None:
# Get status quo values
sq_df = data.df[data.df["arm_name"] ==
experiment.status_quo.name # pyre-ignore
]
sq_df = sq_df.to_dict(orient="list") # pyre-ignore
sq_means = {}
sq_sems = {}
for i, metric in enumerate(sq_df["metric_name"]):
sq_means[metric] = sq_df["mean"][i]
sq_sems[metric] = sq_df["sem"][i]
# Relativize
for name in metric_names:
if np.isnan(sq_sems[name]) or np.isnan(pfr_sems[name]).any():
# Just relativize means
pfr_means[name] = [(mu / sq_means[name] - 1) * 100
for mu in pfr_means[name]]
else:
# Use delta method
pfr_means[name], pfr_sems[name] = relativize(
means_t=pfr_means[name],
sems_t=pfr_sems[name],
mean_c=sq_means[name],
sem_c=sq_sems[name],
as_percent=True,
)
absolute_metrics = []
else:
absolute_metrics = metric_names
objective_thresholds = {}
if experiment.optimization_config.objective_thresholds is not None: # pyre-ignore
for objth in experiment.optimization_config.objective_thresholds:
is_rel = objth.metric.name not in absolute_metrics
if objth.relative != is_rel:
raise ValueError(
f"Objective threshold for {objth.metric.name} has "
f"rel={objth.relative} but was specified here as rel={is_rel}"
)
objective_thresholds[objth.metric.name] = objth.bound
# Construct ParetoFrontResults for each pair
pfr_list = []
param_dicts = [obs.features.parameters for obs in pareto_observations]
arm_names = [obs.arm_name for obs in pareto_observations]
for metric_a, metric_b in combinations(metric_names, 2):
pfr_list.append(
ParetoFrontierResults(
param_dicts=param_dicts,
means=pfr_means,
sems=pfr_sems,
primary_metric=metric_a,
secondary_metric=metric_b,
absolute_metrics=absolute_metrics,
objective_thresholds=objective_thresholds,
arm_names=arm_names,
))
return pfr_list
def testTorchModelBridge(self, mock_init):
torch_dtype = torch.float64
torch_device = torch.device("cpu")
ma = TorchModelBridge(
experiment=None,
search_space=None,
data=None,
model=None,
transforms=[],
torch_dtype=torch.float64,
torch_device=torch.device("cpu"),
)
self.assertEqual(ma.dtype, torch.float64)
self.assertEqual(ma.device, torch.device("cpu"))
# Fit
model = mock.MagicMock(TorchModel, autospec=True, instance=True)
X = np.array([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]])
Y = np.array([[3.0], [4.0]])
var = np.array([[1.0], [2.0]])
ma._model_fit(
model=model,
Xs=[X],
Ys=[Y],
Yvars=[var],
bounds=None,
feature_names=[],
task_features=[],
fidelity_features=[],
)
model_fit_args = model.fit.mock_calls[0][2]
self.assertTrue(
torch.equal(
model_fit_args["Xs"][0],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_fit_args["Ys"][0],
torch.tensor(Y, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_fit_args["Yvars"][0],
torch.tensor(var, dtype=torch_dtype, device=torch_device),
))
# Update
ma._model_update(Xs=[X], Ys=[Y], Yvars=[var])
model_update_args = model.update.mock_calls[0][2]
self.assertTrue(
torch.equal(
model_update_args["Xs"][0],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_update_args["Ys"][0],
torch.tensor(Y, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_update_args["Yvars"][0],
torch.tensor(var, dtype=torch_dtype, device=torch_device),
))
# Predict
model.predict.return_value = (torch.tensor([3.0]), torch.tensor([4.0]))
f, var = ma._model_predict(X)
self.assertTrue(
torch.equal(
model.predict.mock_calls[0][2]["X"],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(np.array_equal(f, np.array([3.0])))
self.assertTrue(np.array_equal(var, np.array([4.0])))
# Gen
model.gen.return_value = (torch.tensor([1.0, 2.0,
3.0]), torch.tensor([1.0]))
X, w = ma._model_gen(
n=3,
bounds=[(0, 1)],
objective_weights=np.array([1.0, 0.0]),
outcome_constraints=None,
linear_constraints=None,
fixed_features={1: 3.0},
pending_observations=[np.array([]),
np.array([1.0, 2.0, 3.0])],
model_gen_options={"option": "yes"},
rounding_func=np.round,
)
gen_args = model.gen.mock_calls[0][2]
self.assertEqual(gen_args["n"], 3)
self.assertEqual(gen_args["bounds"], [(0, 1)])
self.assertTrue(
torch.equal(
gen_args["objective_weights"],
torch.tensor([1.0, 0.0],
dtype=torch_dtype,
device=torch_device),
))
self.assertIsNone(gen_args["outcome_constraints"])
self.assertIsNone(gen_args["linear_constraints"])
self.assertEqual(gen_args["fixed_features"], {1: 3.0})
self.assertTrue(
torch.equal(
gen_args["pending_observations"][0],
torch.tensor([], dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
gen_args["pending_observations"][1],
torch.tensor([1.0, 2.0, 3.0],
dtype=torch_dtype,
device=torch_device),
))
self.assertEqual(gen_args["model_gen_options"], {"option": "yes"})
self.assertTrue(np.array_equal(X, np.array([1.0, 2.0, 3.0])))
self.assertTrue(np.array_equal(w, np.array([1.0])))
# Cross-validate
model.cross_validate.return_value = (torch.tensor([3.0]),
torch.tensor([4.0]))
f, var = ma._model_cross_validate(Xs_train=[X],
Ys_train=[Y],
Yvars_train=[var],
X_test=X)
model_cv_args = model.cross_validate.mock_calls[0][2]
self.assertTrue(
torch.equal(
model_cv_args["Xs_train"][0],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_cv_args["Ys_train"][0],
torch.tensor(Y, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_cv_args["Yvars_train"][0],
torch.tensor(var, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(
torch.equal(
model_cv_args["X_test"],
torch.tensor(X, dtype=torch_dtype, device=torch_device),
))
self.assertTrue(np.array_equal(f, np.array([3.0])))
self.assertTrue(np.array_equal(var, np.array([4.0])))
