def test_construct(self, mock_GP): with self.assertRaises(NotImplementedError): # Base `Model` does not implement `construct_inputs`. Surrogate(botorch_model_class=Model).construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) mock_GP.assert_called_once() call_kwargs = mock_GP.call_args[1] self.assertTrue(torch.equal(call_kwargs["train_X"], self.Xs[0])) self.assertTrue(torch.equal(call_kwargs["train_Y"], self.Ys[0])) self.assertFalse(self.surrogate._constructed_manually) # Check that `model_options` passed to the `Surrogate` constructor are # properly propagated. with patch.object( SingleTaskGP, "construct_inputs", wraps=SingleTaskGP.construct_inputs) as mock_construct_inputs: surrogate = Surrogate( botorch_model_class=self.botorch_model_class, mll_class=self.mll_class, model_options={"some_option": "some_value"}, ) surrogate.construct(self.training_data) mock_construct_inputs.assert_called_with( training_data=self.training_data, some_option="some_value")
def test_construct(self, mock_GP): base_surrogate = Surrogate(botorch_model_class=Model) with self.assertRaisesRegex(TypeError, "Cannot construct an abstract model."): base_surrogate.construct( training_data=self.training_data, fidelity_features=self.fidelity_features, ) self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) mock_GP.assert_called_with(train_X=self.X, train_Y=self.Y)
class AcquisitionSetUp: 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.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 = qMaxValueEntropy 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, } self.optimizer_options = { Keys.NUM_RESTARTS: 40, Keys.RAW_SAMPLES: 1024, Keys.FRAC_RANDOM: 0.2, } self.inequality_constraints = [(torch.tensor([0, 1]), torch.tensor([-1.0, 1.0]), 1)]
class AcquisitionTest(TestCase): 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 tearDown(self): # Avoid polluting the registry for other tests. ACQF_INPUT_CONSTRUCTOR_REGISTRY.pop(DummyACQFClass) @mock.patch(f"{ACQUISITION_PATH}._get_X_pending_and_observed") @mock.patch( f"{ACQUISITION_PATH}.subset_model", return_value=SubsetModelData(None, torch.ones(1), None, None, None), ) @mock.patch(f"{ACQUISITION_PATH}.get_botorch_objective_and_transform") @mock.patch( f"{CURRENT_PATH}.Acquisition.compute_model_dependencies", return_value={"current_value": 1.2}, ) @mock.patch( f"{DummyACQFClass.__module__}.DummyACQFClass.__init__", return_value=None ) 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) @mock.patch(f"{ACQUISITION_PATH}.optimize_acqf") def test_optimize(self, mock_optimize_acqf): self.acquisition.optimize( n=3, search_space_digest=self.search_space_digest, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, rounding_func=self.rounding_func, optimizer_options=self.optimizer_options, ) mock_optimize_acqf.assert_called_with( acq_function=self.acquisition.acqf, bounds=mock.ANY, q=3, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, 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( self.search_space_digest.bounds, dtype=self.acquisition.dtype, device=self.acquisition.device, ).transpose(0, 1) self.assertTrue( torch.equal(mock_optimize_acqf.call_args[1]["bounds"], expected_bounds) ) @mock.patch(f"{ACQUISITION_PATH}.optimize_acqf_discrete") def test_optimize_discrete(self, mock_optimize_acqf_discrete): tkwargs = { "dtype": self.acquisition.dtype, "device": self.acquisition.device, } ssd1 = SearchSpaceDigest( feature_names=["a"], bounds=[(0, 2)], categorical_features=[0], discrete_choices={0: [0, 1, 2]}, ) # check fixed_feature index validation with self.assertRaisesRegex(ValueError, "Invalid fixed_feature index"): self.acquisition.optimize( n=3, search_space_digest=ssd1, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, rounding_func=self.rounding_func, optimizer_options=self.optimizer_options, ) # check this works without any fixed_feature specified self.acquisition.optimize( n=3, search_space_digest=ssd1, inequality_constraints=self.inequality_constraints, fixed_features=None, rounding_func=self.rounding_func, optimizer_options=self.optimizer_options, ) mock_optimize_acqf_discrete.assert_called_with( acq_function=self.acquisition.acqf, q=3, choices=mock.ANY, **self.optimizer_options, ) # can't use assert_called_with on choices due to ambiguous bool comparison expected_choices = torch.tensor([[0], [1], [2]], **tkwargs) self.assertTrue( torch.equal( mock_optimize_acqf_discrete.call_args[1]["choices"], expected_choices ) ) # check with fixed feature ssd2 = SearchSpaceDigest( feature_names=["a", "b"], bounds=[(0, 2), (0, 1)], categorical_features=[0], discrete_choices={0: [0, 1, 2]}, ) self.acquisition.optimize( n=3, search_space_digest=ssd2, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, rounding_func=self.rounding_func, optimizer_options=self.optimizer_options, ) mock_optimize_acqf_discrete.assert_called_with( acq_function=self.acquisition.acqf, q=3, choices=mock.ANY, **self.optimizer_options, ) # can't use assert_called_with on choices due to ambiguous bool comparison expected_choices = torch.tensor([[0, 2.0], [1, 2.0], [2, 2.0]], **tkwargs) self.assertTrue( torch.equal( mock_optimize_acqf_discrete.call_args[1]["choices"], expected_choices ) ) @mock.patch(f"{ACQUISITION_PATH}.optimize_acqf_mixed") 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 ) ) @mock.patch(f"{SURROGATE_PATH}.Surrogate.best_in_sample_point") def test_best_point(self, mock_best_point): self.acquisition.best_point( 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, options=self.options, ) mock_best_point.assert_called_with( 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, options=self.options, ) @mock.patch( f"{DummyACQFClass.__module__}.DummyACQFClass.__call__", return_value=None ) def test_evaluate(self, mock_call): self.acquisition.evaluate(X=self.X) mock_call.assert_called_with(X=self.X) @mock.patch(f"{ACQUISITION_PATH}._get_X_pending_and_observed") 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()))
class MOOAcquisitionTest(TestCase): 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.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, -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.target_fidelities = {2: 1.0} self.options = {} self.acquisition = MOOAcquisition( 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 } @patch(f"{MOO_ACQUISITION_PATH}.get_outcome_constraint_transforms") @patch(f"{ACQUISITION_PATH}._get_X_pending_and_observed") @patch( f"{ACQUISITION_PATH}.subset_model", wraps=acquisition.subset_model, ) @patch(f"{CURRENT_PATH}.MOOAcquisition._get_botorch_objective") @patch(f"{DummyACQFClass.__module__}.DummyACQFClass.__init__", return_value=None) def test_init( self, mock_botorch_acqf_class, mock_get_objective, mock_subset_model, mock_get_X, mock_get_constraints, ): botorch_objective = WeightedMCMultiOutputObjective( weights=self.objective_weights[:2], outcomes=[0, 1]) mock_get_objective.return_value = botorch_objective mock_get_constraints.return_value = self.con_tfs mock_get_X.return_value = (self.pending_observations[0], self.X[:1]) acquisition = MOOAcquisition( 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, ) # Check `_get_X_pending_and_observed` kwargs mock_get_X.assert_called_with( Xs=[ self.training_data.X, self.training_data.X, 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_botorch_acqf_class.reset_mock() self.options[Keys.SUBSET_MODEL] = False acquisition = MOOAcquisition( 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, ) mock_subset_model.assert_not_called() # Check final `acqf` creation 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])) self.assertEqual( ckwargs["ref_point"], (self.objective_thresholds[:2] * self.objective_weights[:2]).tolist(), ) self.assertIsInstance(ckwargs["partitioning"], BoxDecomposition) self.assertIs(ckwargs["constraints"], self.con_tfs) self.assertIsInstance(ckwargs["sampler"], SobolQMCNormalSampler) # qNoisyExpectedImprovement not supported. with self.assertRaisesRegex( UnsupportedError, "Only qExpectedHypervolumeImprovement is currently supported", ): MOOAcquisition( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, objective_thresholds=self.objective_thresholds, botorch_acqf_class=qNoisyExpectedImprovement, 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, ) with self.assertRaisesRegex(ValueError, "Objective Thresholds required"): MOOAcquisition( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, objective_thresholds=None, 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, ) @patch(f"{DummyACQFClass.__module__}.DummyACQFClass.__call__", return_value=None) def test_evaluate(self, mock_call): self.acquisition.evaluate(X=self.X) mock_call.assert_called_with(X=self.X) def test_extract_training_data(self): self.assertEqual( # Base `Surrogate` case. self.acquisition._extract_training_data(surrogate=self.surrogate), self.training_data, ) # `ListSurrogate` case. list_surrogate = ListSurrogate( botorch_submodel_class=self.botorch_model_class) list_surrogate._training_data_per_outcome = {"a": self.training_data} self.assertEqual( self.acquisition._extract_training_data(surrogate=list_surrogate), list_surrogate._training_data_per_outcome, )
class SurrogateTest(TestCase): def setUp(self): self.botorch_model_class = SingleTaskGP self.mll_class = ExactMarginalLogLikelihood self.device = torch.device("cpu") self.dtype = torch.float self.X = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], dtype=self.dtype, device=self.device) self.Y = torch.tensor([[3.0], [4.0]], dtype=self.dtype, device=self.device) self.Yvar = torch.tensor([[0.0], [2.0]], dtype=self.dtype, device=self.device) self.training_data = TrainingData(X=self.X, Y=self.Y, Yvar=self.Yvar) 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.bounds = [(0.0, 1.0), (1.0, 4.0), (2.0, 5.0)] self.task_features = [] self.feature_names = ["x1", "x2", "x3"] self.metric_names = ["y"] self.fidelity_features = [] self.target_fidelities = {1: 1.0} 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 = {} @patch(f"{CURRENT_PATH}.Kernel") @patch(f"{CURRENT_PATH}.Likelihood") def test_init(self, mock_Likelihood, mock_Kernel): self.assertEqual(self.surrogate.botorch_model_class, self.botorch_model_class) self.assertEqual(self.surrogate.mll_class, self.mll_class) with self.assertRaisesRegex(NotImplementedError, "Customizing likelihood"): Surrogate(botorch_model_class=self.botorch_model_class, likelihood=Likelihood()) with self.assertRaisesRegex(NotImplementedError, "Customizing kernel"): Surrogate(botorch_model_class=self.botorch_model_class, kernel_class=Kernel()) def test_model_property(self): with self.assertRaisesRegex( ValueError, "BoTorch `Model` has not yet been constructed."): self.surrogate.model def test_training_data_property(self): with self.assertRaisesRegex( ValueError, "Underlying BoTorch `Model` has not yet received its training_data.", ): self.surrogate.training_data def test_dtype_property(self): self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) self.assertEqual(self.dtype, self.surrogate.dtype) def test_device_property(self): self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) self.assertEqual(self.device, self.surrogate.device) def test_from_BoTorch(self): surrogate = Surrogate.from_BoTorch( self.botorch_model_class(**self.surrogate_kwargs)) self.assertIsInstance(surrogate.model, self.botorch_model_class) self.assertFalse(surrogate._should_reconstruct) @patch(f"{CURRENT_PATH}.SingleTaskGP.__init__", return_value=None) def test_construct(self, mock_GP): base_surrogate = Surrogate(botorch_model_class=Model) with self.assertRaisesRegex(TypeError, "Cannot construct an abstract model."): base_surrogate.construct( training_data=self.training_data, fidelity_features=self.fidelity_features, ) self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) mock_GP.assert_called_with(train_X=self.X, train_Y=self.Y) @patch(f"{CURRENT_PATH}.SingleTaskGP.load_state_dict", return_value=None) @patch(f"{CURRENT_PATH}.ExactMarginalLogLikelihood") @patch(f"{SURROGATE_PATH}.fit_gpytorch_model") def test_fit(self, mock_fit_gpytorch, mock_MLL, mock_state_dict): surrogate = Surrogate( botorch_model_class=self.botorch_model_class, mll_class=ExactMarginalLogLikelihood, ) # 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, bounds=self.bounds, task_features=self.task_features, feature_names=self.feature_names, metric_names=self.metric_names, fidelity_features=self.fidelity_features, target_fidelities=self.target_fidelities, refit=self.refit, ) 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 = {} surrogate.fit( training_data=self.training_data, bounds=self.bounds, task_features=self.task_features, feature_names=self.feature_names, metric_names=self.metric_names, fidelity_features=self.fidelity_features, target_fidelities=self.target_fidelities, refit=False, state_dict=state_dict, ) mock_state_dict.assert_called_once() mock_MLL.assert_not_called() mock_fit_gpytorch.assert_not_called() @patch(f"{SURROGATE_PATH}.predict_from_model") def test_predict(self, mock_predict): self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) self.surrogate.predict(X=self.X) mock_predict.assert_called_with(model=self.surrogate.model, X=self.X) def test_best_in_sample_point(self): self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) # `best_in_sample_point` requires `objective_weights` with patch(f"{SURROGATE_PATH}.best_in_sample_point", return_value=None) as mock_best_in_sample: with self.assertRaisesRegex(ValueError, "Could not obtain"): self.surrogate.best_in_sample_point(bounds=self.bounds, objective_weights=None) with patch(f"{SURROGATE_PATH}.best_in_sample_point", return_value=(self.X, 0.0)) as mock_best_in_sample: best_point, observed_value = self.surrogate.best_in_sample_point( bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) mock_best_in_sample.assert_called_with( Xs=[self.training_data.X], model=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) @patch(f"{ACQUISITION_PATH}.Acquisition.__init__", return_value=None) @patch( f"{ACQUISITION_PATH}.Acquisition.optimize", return_value=([torch.tensor([0.0])], [torch.tensor([1.0])]), ) @patch( f"{SURROGATE_PATH}.pick_best_out_of_sample_point_acqf_class", return_value=(qSimpleRegret, { Keys.SAMPLER: SobolQMCNormalSampler }), ) def test_best_out_of_sample_point(self, mock_best_point_util, mock_acqf_optimize, mock_acqf_init): self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) # currently cannot use function with fixed features with self.assertRaisesRegex(NotImplementedError, "Fixed features"): self.surrogate.best_out_of_sample_point( bounds=self.bounds, objective_weights=self.objective_weights, fixed_features=self.fixed_features, ) candidate, acqf_value = self.surrogate.best_out_of_sample_point( bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fidelity_features=self.fidelity_features, target_fidelities=self.target_fidelities, options=self.options, ) mock_acqf_init.assert_called_with( surrogate=self.surrogate, botorch_acqf_class=qSimpleRegret, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=None, target_fidelities=self.target_fidelities, options={Keys.SAMPLER: SobolQMCNormalSampler}, ) self.assertTrue(torch.equal(candidate, torch.tensor([0.0]))) self.assertTrue(torch.equal(acqf_value, torch.tensor([1.0]))) @patch(f"{SURROGATE_PATH}.Surrogate.fit") def test_update(self, mock_fit): self.surrogate.construct(training_data=self.training_data, fidelity_features=self.fidelity_features) # Call `fit` by default self.surrogate.update( training_data=self.training_data, bounds=self.bounds, task_features=self.task_features, feature_names=self.feature_names, metric_names=self.metric_names, fidelity_features=self.fidelity_features, refit=self.refit, ) mock_fit.assert_called_with( training_data=self.training_data, bounds=self.bounds, task_features=self.task_features, feature_names=self.feature_names, metric_names=self.metric_names, fidelity_features=self.fidelity_features, candidate_metadata=None, state_dict=self.surrogate.model.state_dict, refit=self.refit, ) # If should not be reconstructed, raise Error self.surrogate._should_reconstruct = False with self.assertRaisesRegex( NotImplementedError, ".* models that should not be re-constructed"): self.surrogate.update( training_data=self.training_data, bounds=self.bounds, task_features=self.task_features, feature_names=self.feature_names, metric_names=self.metric_names, fidelity_features=self.fidelity_features, refit=self.refit, )
class AcquisitionTest(TestCase): 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.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, 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} @patch( f"{ACQUISITION_PATH}._get_X_pending_and_observed", return_value=(torch.tensor([2.0]), torch.tensor([3.0])), ) @patch(f"{ACQUISITION_PATH}.subset_model", return_value=(None, None, None, None)) @patch(f"{ACQUISITION_PATH}.get_botorch_objective") @patch( f"{CURRENT_PATH}.Acquisition.compute_model_dependencies", return_value={"current_value": 1.2}, ) @patch(f"{CURRENT_PATH}.Acquisition.compute_data_dependencies", return_value={}) @patch(f"{DummyACQFClass.__module__}.DummyACQFClass.__init__", return_value=None) 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, ) @patch(f"{ACQUISITION_PATH}.optimize_acqf") def test_optimize(self, mock_optimize_acqf): self.acquisition.optimize( bounds=self.bounds, n=3, optimizer_class=None, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, rounding_func=self.rounding_func, optimizer_options=self.optimizer_options, ) mock_optimize_acqf.assert_called_with( self.acquisition.acqf, bounds=self.bounds, q=3, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, post_processing_func=self.rounding_func, **self.optimizer_options, ) @patch(f"{SURROGATE_PATH}.Surrogate.best_in_sample_point") def test_best_point(self, mock_best_point): self.acquisition.best_point( bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, target_fidelities=self.target_fidelities, options=self.options, ) mock_best_point.assert_called_with( bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) @patch(f"{DummyACQFClass.__module__}.DummyACQFClass.__call__", return_value=None) def test_evaluate(self, mock_call): self.acquisition.evaluate(X=self.X) mock_call.assert_called_with(X=self.X) def test_extract_training_data(self): self.assertEqual( # Base `Surrogate` case. self.acquisition._extract_training_data(surrogate=self.surrogate), self.training_data, ) # `ListSurrogate` case. list_surrogate = ListSurrogate(botorch_submodel_class=self.botorch_model_class) list_surrogate._training_data_per_outcome = {"a": self.training_data} self.assertEqual( self.acquisition._extract_training_data(surrogate=list_surrogate), list_surrogate._training_data_per_outcome, )
class SurrogateTest(TestCase): 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 = {} @patch(f"{CURRENT_PATH}.Kernel") @patch(f"{CURRENT_PATH}.Likelihood") def test_init(self, mock_Likelihood, mock_Kernel): self.assertEqual(self.surrogate.botorch_model_class, self.botorch_model_class) self.assertEqual(self.surrogate.mll_class, self.mll_class) with self.assertRaisesRegex(NotImplementedError, "Customizing likelihood"): Surrogate(botorch_model_class=self.botorch_model_class, likelihood=Likelihood()) with self.assertRaisesRegex(NotImplementedError, "Customizing kernel"): Surrogate(botorch_model_class=self.botorch_model_class, kernel_class=Kernel()) @patch(f"{SURROGATE_PATH}.fit_gpytorch_model") def test_mll_options(self, _): mock_mll = MagicMock(self.mll_class) surrogate = Surrogate( botorch_model_class=self.botorch_model_class, mll_class=mock_mll, mll_options={"some_option": "some_value"}, ) surrogate.fit( training_data=self.training_data, search_space_digest=self.search_space_digest, metric_names=self.metric_names, refit=self.refit, ) self.assertEqual(mock_mll.call_args[1]["some_option"], "some_value") def test_model_property(self): with self.assertRaisesRegex( ValueError, "BoTorch `Model` has not yet been constructed."): self.surrogate.model def test_training_data_property(self): with self.assertRaisesRegex( ValueError, "Underlying BoTorch `Model` has not yet received its training_data.", ): self.surrogate.training_data def test_dtype_property(self): self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) self.assertEqual(self.dtype, self.surrogate.dtype) def test_device_property(self): self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) self.assertEqual(self.device, self.surrogate.device) def test_from_botorch(self): surrogate = Surrogate.from_botorch( self.botorch_model_class(**self.surrogate_kwargs)) self.assertIsInstance(surrogate.model, self.botorch_model_class) self.assertTrue(surrogate._constructed_manually) @patch(f"{CURRENT_PATH}.SingleTaskGP.__init__", return_value=None) def test_construct(self, mock_GP): with self.assertRaises(NotImplementedError): # Base `Model` does not implement `construct_inputs`. Surrogate(botorch_model_class=Model).construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) mock_GP.assert_called_once() call_kwargs = mock_GP.call_args[1] self.assertTrue(torch.equal(call_kwargs["train_X"], self.Xs[0])) self.assertTrue(torch.equal(call_kwargs["train_Y"], self.Ys[0])) self.assertFalse(self.surrogate._constructed_manually) # Check that `model_options` passed to the `Surrogate` constructor are # properly propagated. with patch.object( SingleTaskGP, "construct_inputs", wraps=SingleTaskGP.construct_inputs) as mock_construct_inputs: surrogate = Surrogate( botorch_model_class=self.botorch_model_class, mll_class=self.mll_class, model_options={"some_option": "some_value"}, ) surrogate.construct(self.training_data) mock_construct_inputs.assert_called_with( training_data=self.training_data, some_option="some_value") @patch(f"{CURRENT_PATH}.SingleTaskGP.load_state_dict", return_value=None) @patch(f"{CURRENT_PATH}.ExactMarginalLogLikelihood") @patch(f"{SURROGATE_PATH}.fit_gpytorch_model") def test_fit(self, mock_fit_gpytorch, mock_MLL, mock_state_dict): surrogate = Surrogate( botorch_model_class=self.botorch_model_class, mll_class=ExactMarginalLogLikelihood, ) # 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=self.search_space_digest, metric_names=self.metric_names, refit=self.refit, ) # Check that training data is correctly passed through to the # BoTorch `Model`. self.assertTrue( torch.equal( surrogate.model.train_inputs[0], self.surrogate_kwargs.get("train_X"), )) self.assertTrue( torch.equal( surrogate.model.train_targets, self.surrogate_kwargs.get("train_Y").squeeze(1), )) 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=self.search_space_digest, metric_names=self.metric_names, refit=False, state_dict=state_dict, ) mock_state_dict.assert_called_once() mock_MLL.assert_not_called() mock_fit_gpytorch.assert_not_called() @patch(f"{SURROGATE_PATH}.predict_from_model") def test_predict(self, mock_predict): self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) self.surrogate.predict(X=self.Xs[0]) mock_predict.assert_called_with(model=self.surrogate.model, X=self.Xs[0]) def test_best_in_sample_point(self): self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) # `best_in_sample_point` requires `objective_weights` with patch(f"{SURROGATE_PATH}.best_in_sample_point", return_value=None) as mock_best_in_sample: with self.assertRaisesRegex(ValueError, "Could not obtain"): self.surrogate.best_in_sample_point( search_space_digest=self.search_space_digest, objective_weights=None) with patch(f"{SURROGATE_PATH}.best_in_sample_point", return_value=(self.Xs[0], 0.0)) as mock_best_in_sample: best_point, observed_value = self.surrogate.best_in_sample_point( 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, options=self.options, ) mock_best_in_sample.assert_called_with( Xs=[self.training_data.X], model=self.surrogate, bounds=self.search_space_digest.bounds, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) @patch(f"{ACQUISITION_PATH}.Acquisition.__init__", return_value=None) @patch( f"{ACQUISITION_PATH}.Acquisition.optimize", return_value=([torch.tensor([0.0])], [torch.tensor([1.0])]), ) @patch( f"{SURROGATE_PATH}.pick_best_out_of_sample_point_acqf_class", return_value=(qSimpleRegret, { Keys.SAMPLER: SobolQMCNormalSampler }), ) def test_best_out_of_sample_point(self, mock_best_point_util, mock_acqf_optimize, mock_acqf_init): self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) # currently cannot use function with fixed features with self.assertRaisesRegex(NotImplementedError, "Fixed features"): self.surrogate.best_out_of_sample_point( search_space_digest=self.search_space_digest, objective_weights=self.objective_weights, fixed_features=self.fixed_features, ) candidate, acqf_value = self.surrogate.best_out_of_sample_point( search_space_digest=self.search_space_digest, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, options=self.options, ) mock_acqf_init.assert_called_with( surrogate=self.surrogate, botorch_acqf_class=qSimpleRegret, search_space_digest=self.search_space_digest, objective_weights=self.objective_weights, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=None, options={Keys.SAMPLER: SobolQMCNormalSampler}, ) self.assertTrue(torch.equal(candidate, torch.tensor([0.0]))) self.assertTrue(torch.equal(acqf_value, torch.tensor([1.0]))) @patch(f"{CURRENT_PATH}.SingleTaskGP.load_state_dict", return_value=None) @patch(f"{CURRENT_PATH}.ExactMarginalLogLikelihood") @patch(f"{SURROGATE_PATH}.fit_gpytorch_model") def test_update(self, mock_fit_gpytorch, mock_MLL, mock_state_dict): self.surrogate.construct( training_data=self.training_data, fidelity_features=self.search_space_digest.fidelity_features, ) # Check that correct arguments are passed to `fit`. with patch(f"{SURROGATE_PATH}.Surrogate.fit") as mock_fit: # Call `fit` by default self.surrogate.update( training_data=self.training_data, search_space_digest=self.search_space_digest, metric_names=self.metric_names, refit=self.refit, state_dict={"key": "val"}, ) mock_fit.assert_called_with( training_data=self.training_data, search_space_digest=self.search_space_digest, metric_names=self.metric_names, candidate_metadata=None, refit=self.refit, state_dict={"key": "val"}, ) # Check that the training data is correctly passed through to the # BoTorch `Model`. Xs, Ys, Yvars, bounds, _, _, _ = get_torch_test_data(dtype=self.dtype, offset=1.0) training_data = TrainingData.from_block_design(X=Xs[0], Y=Ys[0], Yvar=Yvars[0]) surrogate_kwargs = self.botorch_model_class.construct_inputs( training_data) self.surrogate.update( training_data=training_data, search_space_digest=self.search_space_digest, metric_names=self.metric_names, refit=self.refit, state_dict={"key": "val"}, ) self.assertTrue( torch.equal( self.surrogate.model.train_inputs[0], surrogate_kwargs.get("train_X"), )) self.assertTrue( torch.equal( self.surrogate.model.train_targets, surrogate_kwargs.get("train_Y").squeeze(1), )) # If should not be reconstructed, check that error is raised. self.surrogate._constructed_manually = True with self.assertRaisesRegex(NotImplementedError, ".* constructed manually"): self.surrogate.update( training_data=self.training_data, search_space_digest=self.search_space_digest, metric_names=self.metric_names, refit=self.refit, ) def test_serialize_attributes_as_kwargs(self): expected = self.surrogate.__dict__ self.assertEqual(self.surrogate._serialize_attributes_as_kwargs(), expected)
class MultiFidelityAcquisitionTest(TestCase): 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.acquisition_options = {Keys.NUM_FANTASIES: 64} self.bounds = [(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)] self.objective_weights = torch.tensor([1.0]) self.target_fidelities = {2: 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, } with patch(f"{MFKG_PATH}.__init__", return_value=None): # We don't actually need to instantiate the BoTorch acqf in these tests. self.acquisition = MultiFidelityAcquisition( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, botorch_acqf_class=qMultiFidelityKnowledgeGradient, target_fidelities=self.target_fidelities, ) @patch(f"{ACQUISITION_PATH}.Acquisition.__init__", return_value=None) @patch(f"{ACQUISITION_PATH}.Acquisition.optimize") def test_optimize(self, mock_Acquisition_optimize, mock_Acquisition_init): # `self.acquisition.optimize()` should call `Acquisition.optimize()` # once. self.acquisition.optimize(bounds=self.bounds, n=1) mock_Acquisition_optimize.assert_called_once() @patch(f"{ACQUISITION_PATH}.Acquisition.compute_model_dependencies", return_value={}) @patch(f"{MULTI_FIDELITY_PATH}.AffineFidelityCostModel", return_value="cost_model") @patch(f"{MULTI_FIDELITY_PATH}.InverseCostWeightedUtility", return_value=None) @patch(f"{MULTI_FIDELITY_PATH}.project_to_target_fidelity", return_value=None) @patch(f"{MULTI_FIDELITY_PATH}.expand_trace_observations", return_value=None) def test_compute_model_dependencies( self, mock_expand, mock_project, mock_inverse_utility, mock_affine_model, mock_Acquisition_compute, ): # Raise Error if `fidelity_weights` and `target_fidelities` do # not align. with self.assertRaisesRegex(RuntimeError, "Must provide the same indices"): self.acquisition.compute_model_dependencies( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, target_fidelities={1: 5.0}, pending_observations=self.pending_observations, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) # Make sure `fidelity_weights` are set when they are not passed in. self.acquisition.compute_model_dependencies( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, target_fidelities={ 2: 5.0, 3: 5.0 }, pending_observations=self.pending_observations, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options={ Keys.COST_INTERCEPT: 1.0, Keys.NUM_TRACE_OBSERVATIONS: 0 }, ) mock_affine_model.assert_called_with(fidelity_weights={ 2: 1.0, 3: 1.0 }, fixed_cost=1.0) # Usual case. dependencies = self.acquisition.compute_model_dependencies( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, target_fidelities=self.target_fidelities, pending_observations=self.pending_observations, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) mock_Acquisition_compute.assert_called_with( surrogate=self.surrogate, bounds=self.bounds, objective_weights=self.objective_weights, target_fidelities=self.target_fidelities, pending_observations=self.pending_observations, outcome_constraints=self.outcome_constraints, linear_constraints=self.linear_constraints, fixed_features=self.fixed_features, options=self.options, ) mock_affine_model.assert_called_with( fidelity_weights=self.options[Keys.FIDELITY_WEIGHTS], fixed_cost=self.options[Keys.COST_INTERCEPT], ) mock_inverse_utility.assert_called_with(cost_model="cost_model") self.assertTrue(Keys.COST_AWARE_UTILITY in dependencies) self.assertTrue(Keys.PROJECT in dependencies) self.assertTrue(Keys.EXPAND in dependencies) # Check that `project` and `expand` are defined correctly. project = dependencies.get(Keys.PROJECT) project(torch.tensor([1.0])) mock_project.assert_called_with( X=torch.tensor([1.0]), target_fidelities=self.target_fidelities) expand = dependencies.get(Keys.EXPAND) expand(torch.tensor([1.0])) mock_expand.assert_called_with( X=torch.tensor([1.0]), fidelity_dims=sorted(self.target_fidelities), num_trace_obs=self.options.get(Keys.NUM_TRACE_OBSERVATIONS), )
class AcquisitionTest(TestCase): 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.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( 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} @mock.patch(f"{ACQUISITION_PATH}._get_X_pending_and_observed") @mock.patch( f"{ACQUISITION_PATH}.subset_model", return_value=(None, None, None, None) ) @mock.patch(f"{ACQUISITION_PATH}.get_botorch_objective") @mock.patch( f"{CURRENT_PATH}.Acquisition.compute_model_dependencies", return_value={"current_value": 1.2}, ) @mock.patch( f"{DummyACQFClass.__module__}.DummyACQFClass.__init__", return_value=None ) 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, search_space_digest=self.search_space_digest, 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, 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, ) # 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( 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, 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` 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) @mock.patch(f"{ACQUISITION_PATH}.optimize_acqf") def test_optimize(self, mock_optimize_acqf): self.acquisition.optimize( n=3, search_space_digest=self.search_space_digest, optimizer_class=None, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, rounding_func=self.rounding_func, optimizer_options=self.optimizer_options, ) mock_optimize_acqf.assert_called_with( acq_function=self.acquisition.acqf, bounds=mock.ANY, q=3, inequality_constraints=self.inequality_constraints, fixed_features=self.fixed_features, 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( self.search_space_digest.bounds, dtype=self.acquisition.dtype, device=self.acquisition.device, ).transpose(0, 1) self.assertTrue( torch.equal(mock_optimize_acqf.call_args[1]["bounds"], expected_bounds) ) @mock.patch(f"{SURROGATE_PATH}.Surrogate.best_in_sample_point") def test_best_point(self, mock_best_point): self.acquisition.best_point( 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, options=self.options, ) mock_best_point.assert_called_with( 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, options=self.options, ) @mock.patch( f"{DummyACQFClass.__module__}.DummyACQFClass.__call__", return_value=None ) def test_evaluate(self, mock_call): self.acquisition.evaluate(X=self.X) mock_call.assert_called_with(X=self.X) def test_extract_training_data(self): self.assertEqual( # Base `Surrogate` case. self.acquisition._extract_training_data(surrogate=self.surrogate), self.training_data, ) # `ListSurrogate` case. list_surrogate = ListSurrogate(botorch_submodel_class=self.botorch_model_class) list_surrogate._training_data_per_outcome = {"a": self.training_data} self.assertEqual( self.acquisition._extract_training_data(surrogate=list_surrogate), list_surrogate._training_data_per_outcome, )