def testUnwrapObservationData(self):
observation_data = [get_observation1().data, get_observation2().data]
f, cov = unwrap_observation_data(observation_data)
self.assertEqual(f["a"], [2.0, 2.0])
self.assertEqual(f["b"], [4.0, 1.0])
self.assertEqual(cov["a"]["a"], [1.0, 2.0])
self.assertEqual(cov["b"]["b"], [4.0, 5.0])
self.assertEqual(cov["a"]["b"], [2.0, 3.0])
self.assertEqual(cov["b"]["a"], [3.0, 4.0])
# Check that errors if metric mismatch
od3 = ObservationData(metric_names=["a"],
means=np.array([2.0]),
covariance=np.array([[4.0]]))
with self.assertRaises(ValueError):
unwrap_observation_data(observation_data + [od3])
def transform_optimization_config(
self,
optimization_config: OptimizationConfig,
modelbridge: Optional["modelbridge_module.base.ModelBridge"],
fixed_features: ObservationFeatures,
) -> OptimizationConfig:
use_raw_sq = self.config.get("use_raw_status_quo", False)
has_relative_constraint = any(
c.relative for c in optimization_config.all_constraints
)
if not has_relative_constraint:
return optimization_config
# Else, we have at least one relative constraint.
# Estimate the value at the status quo.
if modelbridge is None:
raise ValueError("ModelBridge not supplied to transform.")
if modelbridge.status_quo is None:
raise ValueError(
"Optimization config has relative constraint, but model was "
"not fit with status quo."
)
try:
# pyre-fixme[16]: `Optional` has no attribute `features`.
f, _ = modelbridge.predict([modelbridge.status_quo.features])
except Exception:
# Check if it is out-of-design.
if use_raw_sq or not modelbridge.model_space.check_membership(
modelbridge.status_quo.features.parameters
):
# Out-of-design: use the raw observation
sq_data = ivw_metric_merge(
# pyre-fixme[16]: `Optional` has no attribute `data`.
obsd=modelbridge.status_quo.data,
conflicting_noiseless="raise",
)
f, _ = unwrap_observation_data([sq_data])
else:
# Should have worked.
raise # pragma: no cover
# Plug in the status quo value to each relative constraint.
for c in optimization_config.all_constraints:
if c.relative:
if isinstance(c, ScalarizedOutcomeConstraint):
# The sq_val of scalarized outcome is the weighted
# sum of its component metrics
sq_val = np.sum(
[
c.weights[i] * f[metric.name][0]
for i, metric in enumerate(c.metrics)
]
)
else:
sq_val = f[c.metric.name][0]
c.bound = (1 + c.bound / 100.0) * sq_val
c.relative = False
return optimization_config
def transform_optimization_config(
self,
optimization_config: OptimizationConfig,
modelbridge: Optional["modelbridge_module.base.ModelBridge"],
fixed_features: ObservationFeatures,
) -> OptimizationConfig:
has_relative_constraint = any(
c.relative for c in optimization_config.outcome_constraints
)
if not has_relative_constraint:
return optimization_config
# Else, we have at least one relative constraint.
# Estimate the value at the status quo.
if modelbridge is None:
raise ValueError("ModelBridge not supplied to transform.")
if modelbridge.status_quo is None:
raise ValueError(
"Optimization config has relative constraint, but model was "
"not fit with status quo."
)
try:
# pyre-fixme[16]: `Optional` has no attribute `features`.
f, _ = modelbridge.predict([modelbridge.status_quo.features])
except Exception:
# Check if it is out-of-design.
if not modelbridge.model_space.check_membership(
modelbridge.status_quo.features.parameters
):
# Out-of-design: use the raw observation
sq_data = ivw_metric_merge(
# pyre-fixme[16]: `Optional` has no attribute `data`.
obsd=modelbridge.status_quo.data,
conflicting_noiseless="raise",
)
f, _ = unwrap_observation_data([sq_data])
else:
# Should have worked.
raise # pragma: no cover
# Plug in the status quo value to each relative constraint.
for c in optimization_config.outcome_constraints:
if c.relative:
# Compute new bound.
c.bound = (1 + c.bound / 100.0) * f[c.metric.name][0]
c.relative = False
return optimization_config
def testModelBridge(self, mock_fit, mock_gen_arms,
mock_observations_from_data):
# Test that on init transforms are stored and applied in the correct order
transforms = [transform_1, transform_2]
exp = get_experiment_for_value()
ss = get_search_space_for_value()
modelbridge = ModelBridge(ss, 0, transforms, exp, 0)
self.assertEqual(list(modelbridge.transforms.keys()),
["transform_1", "transform_2"])
fit_args = mock_fit.mock_calls[0][2]
self.assertTrue(
fit_args["search_space"] == get_search_space_for_value(8.0))
self.assertTrue(fit_args["observation_features"] == [])
self.assertTrue(fit_args["observation_data"] == [])
self.assertTrue(mock_observations_from_data.called)
# Test prediction on out of design features.
modelbridge._predict = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._predict",
autospec=True,
side_effect=ValueError("Out of Design"),
)
# This point is in design, and thus failures in predict are legitimate.
with mock.patch.object(ModelBridge,
"model_space",
return_value=get_search_space_for_range_values):
with self.assertRaises(ValueError):
modelbridge.predict([get_observation2().features])
# This point is out of design, and not in training data.
with self.assertRaises(ValueError):
modelbridge.predict([get_observation_status_quo0().features])
# Now it's in the training data.
with mock.patch.object(
ModelBridge,
"get_training_data",
return_value=[get_observation_status_quo0()],
):
# Return raw training value.
self.assertEqual(
modelbridge.predict([get_observation_status_quo0().features]),
unwrap_observation_data([get_observation_status_quo0().data]),
)
# Test that transforms are applied correctly on predict
modelbridge._predict = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._predict",
autospec=True,
return_value=[get_observation2trans().data],
)
modelbridge.predict([get_observation2().features])
# Observation features sent to _predict are un-transformed afterwards
modelbridge._predict.assert_called_with([get_observation2().features])
# Check that _single_predict is equivalent here.
modelbridge._single_predict([get_observation2().features])
# Observation features sent to _predict are un-transformed afterwards
modelbridge._predict.assert_called_with([get_observation2().features])
# Test transforms applied on gen
modelbridge._gen = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._gen",
autospec=True,
return_value=([get_observation1trans().features], [2], None, {}),
)
oc = OptimizationConfig(objective=Objective(metric=Metric(
name="test_metric")))
modelbridge._set_kwargs_to_save(model_key="TestModel",
model_kwargs={},
bridge_kwargs={})
gr = modelbridge.gen(
n=1,
search_space=get_search_space_for_value(),
optimization_config=oc,
pending_observations={"a": [get_observation2().features]},
fixed_features=ObservationFeatures({"x": 5}),
)
self.assertEqual(gr._model_key, "TestModel")
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace(
[FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=oc,
pending_observations={"a": [get_observation2trans().features]},
fixed_features=ObservationFeatures({"x": 36}),
model_gen_options=None,
)
mock_gen_arms.assert_called_with(
arms_by_signature={},
observation_features=[get_observation1().features])
# Gen with no pending observations and no fixed features
modelbridge.gen(n=1,
search_space=get_search_space_for_value(),
optimization_config=None)
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace(
[FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=None,
pending_observations={},
fixed_features=ObservationFeatures({}),
model_gen_options=None,
)
# Gen with multi-objective optimization config.
oc2 = OptimizationConfig(objective=ScalarizedObjective(
metrics=[Metric(name="test_metric"),
Metric(name="test_metric_2")]))
modelbridge.gen(n=1,
search_space=get_search_space_for_value(),
optimization_config=oc2)
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace(
[FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=oc2,
pending_observations={},
fixed_features=ObservationFeatures({}),
model_gen_options=None,
)
# Test transforms applied on cross_validate
modelbridge._cross_validate = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._cross_validate",
autospec=True,
return_value=[get_observation1trans().data],
)
cv_training_data = [get_observation2()]
cv_test_points = [get_observation1().features]
cv_predictions = modelbridge.cross_validate(
cv_training_data=cv_training_data, cv_test_points=cv_test_points)
modelbridge._cross_validate.assert_called_with(
obs_feats=[get_observation2trans().features],
obs_data=[get_observation2trans().data],
cv_test_points=[get_observation1().features
], # untransformed after
)
self.assertTrue(cv_predictions == [get_observation1().data])
# Test stored training data
obs = modelbridge.get_training_data()
self.assertTrue(obs == [get_observation1(), get_observation2()])
self.assertEqual(modelbridge.metric_names, {"a", "b"})
self.assertIsNone(modelbridge.status_quo)
self.assertTrue(
modelbridge.model_space == get_search_space_for_value())
self.assertEqual(modelbridge.training_in_design, [False, False])
with self.assertRaises(ValueError):
modelbridge.training_in_design = [True, True, False]
with self.assertRaises(ValueError):
modelbridge.training_in_design = [True, True, False]
# Test feature_importances
with self.assertRaises(NotImplementedError):
modelbridge.feature_importances("a")
def testModelBridge(self, mock_fit, mock_gen_arms, mock_observations_from_data):
# Test that on init transforms are stored and applied in the correct order
transforms = [t1, t2]
exp = get_experiment()
modelbridge = ModelBridge(search_space_for_value(), 0, transforms, exp, 0)
self.assertEqual(list(modelbridge.transforms.keys()), ["t1", "t2"])
fit_args = mock_fit.mock_calls[0][2]
self.assertTrue(fit_args["search_space"] == search_space_for_value(8.0))
self.assertTrue(
fit_args["observation_features"]
== [observation1trans().features, observation2trans().features]
)
self.assertTrue(
fit_args["observation_data"]
== [observation1trans().data, observation2trans().data]
)
self.assertTrue(mock_observations_from_data.called)
# Test that transforms are applied correctly on predict
modelbridge._predict = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._predict",
autospec=True,
return_value=[observation2trans().data],
)
modelbridge.predict([observation2().features])
# Observation features sent to _predict are un-transformed afterwards
modelbridge._predict.assert_called_with([observation2().features])
# Test transforms applied on gen
modelbridge._gen = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._gen",
autospec=True,
return_value=([observation1trans().features], [2], None),
)
oc = OptimizationConfig(objective=Objective(metric=Metric(name="test_metric")))
modelbridge._set_kwargs_to_save(
model_key="TestModel", model_kwargs={}, bridge_kwargs={}
)
gr = modelbridge.gen(
n=1,
search_space=search_space_for_value(),
optimization_config=oc,
pending_observations={"a": [observation2().features]},
fixed_features=ObservationFeatures({"x": 5}),
)
self.assertEqual(gr._model_key, "TestModel")
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=oc,
pending_observations={"a": [observation2trans().features]},
fixed_features=ObservationFeatures({"x": 36}),
model_gen_options=None,
)
mock_gen_arms.assert_called_with(
arms_by_signature={}, observation_features=[observation1().features]
)
# Gen with no pending observations and no fixed features
modelbridge.gen(
n=1, search_space=search_space_for_value(), optimization_config=None
)
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=None,
pending_observations={},
fixed_features=ObservationFeatures({}),
model_gen_options=None,
)
# Gen with multi-objective optimization config.
oc2 = OptimizationConfig(
objective=ScalarizedObjective(
metrics=[Metric(name="test_metric"), Metric(name="test_metric_2")]
)
)
modelbridge.gen(
n=1, search_space=search_space_for_value(), optimization_config=oc2
)
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=oc2,
pending_observations={},
fixed_features=ObservationFeatures({}),
model_gen_options=None,
)
# Test transforms applied on cross_validate
modelbridge._cross_validate = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._cross_validate",
autospec=True,
return_value=[observation1trans().data],
)
cv_training_data = [observation2()]
cv_test_points = [observation1().features]
cv_predictions = modelbridge.cross_validate(
cv_training_data=cv_training_data, cv_test_points=cv_test_points
)
modelbridge._cross_validate.assert_called_with(
obs_feats=[observation2trans().features],
obs_data=[observation2trans().data],
cv_test_points=[observation1().features], # untransformed after
)
self.assertTrue(cv_predictions == [observation1().data])
# Test stored training data
obs = modelbridge.get_training_data()
self.assertTrue(obs == [observation1(), observation2()])
self.assertEqual(modelbridge.metric_names, {"a", "b"})
self.assertIsNone(modelbridge.status_quo)
self.assertTrue(modelbridge.model_space == search_space_for_value())
self.assertEqual(modelbridge.training_in_design, [True, True])
modelbridge.training_in_design = [True, False]
with self.assertRaises(ValueError):
modelbridge.training_in_design = [True, True, False]
ood_obs = modelbridge.out_of_design_data()
self.assertTrue(ood_obs == unwrap_observation_data([observation2().data]))
