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 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"])
