def test_sample_points_around_best(self): tkwargs = {"device": self.device} _bounds = torch.ones(2, 2) _bounds[1] = 2 for dtype in (torch.float, torch.double): tkwargs["dtype"] = dtype bounds = _bounds.to(**tkwargs) X_train = 1 + torch.rand(20, 2, **tkwargs) model = MockModel( MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)) ) # test NEI with X_baseline acqf = qNoisyExpectedImprovement(model, X_baseline=X_train) with mock.patch( "botorch.optim.initializers.sample_perturbed_subset_dims" ) as mock_subset_dims: X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) mock_subset_dims.assert_not_called() self.assertTrue(X_rnd.shape, torch.Size([4, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test model that returns a batched mean model = MockModel( MockPosterior( mean=(2 * X_train + 1).sum(dim=-1, keepdim=True).unsqueeze(0) ) ) acqf = qNoisyExpectedImprovement(model, X_baseline=X_train) X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertTrue(X_rnd.shape, torch.Size([4, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test EI without X_baseline acqf = qExpectedImprovement(model, best_f=0.0) with warnings.catch_warnings(record=True) as w, settings.debug(True): X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) # set train inputs model.train_inputs = (X_train,) X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertTrue(X_rnd.shape, torch.Size([4, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test an acquisition function that has objective=None # and maximize=False pm = PosteriorMean(model, maximize=False) self.assertIsNone(pm.objective) self.assertFalse(pm.maximize) X_rnd = sample_points_around_best( acq_function=pm, n_discrete_points=4, sigma=0, bounds=bounds, best_pct=1e-8, # ensures that we only use best value ) idx = (-model.posterior(X_train).mean).argmax() self.assertTrue((X_rnd == X_train[idx : idx + 1]).all(dim=-1).all()) # test acquisition function that has no model ff = FixedFeatureAcquisitionFunction(pm, d=2, columns=[0], values=[0]) # set X_baseline for testing purposes ff.X_baseline = X_train with warnings.catch_warnings(record=True) as w, settings.debug(True): X_rnd = sample_points_around_best( acq_function=ff, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) # test constraints with NEHVI constraints = [lambda Y: Y[..., 0]] ref_point = torch.zeros(2, **tkwargs) # test cases when there are and are not any feasible points for any_feas in (True, False): Y_train = torch.stack( [ torch.linspace(-0.5, 0.5, X_train.shape[0], **tkwargs) if any_feas else torch.ones(X_train.shape[0], **tkwargs), X_train.sum(dim=-1), ], dim=-1, ) moo_model = MockModel(MockPosterior(mean=Y_train, samples=Y_train)) acqf = qNoisyExpectedHypervolumeImprovement( moo_model, ref_point=ref_point, X_baseline=X_train, constraints=constraints, cache_root=False, ) X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=0.0, bounds=bounds, ) self.assertTrue(X_rnd.shape, torch.Size([4, 2])) # this should be true since sigma=0 # and we should only be returning feasible points violation = constraints[0](Y_train) neg_violation = -violation.clamp_min(0.0) feas = neg_violation == 0 eq_mask = (X_train.unsqueeze(1) == X_rnd.unsqueeze(0)).all(dim=-1) if feas.any(): # determine # create n_train x n_rnd tensor of booleans eq_mask = (X_train.unsqueeze(1) == X_rnd.unsqueeze(0)).all(dim=-1) # check that all X_rnd correspond to feasible points self.assertEqual(eq_mask[feas].sum(), 4) else: idcs = torch.topk(neg_violation, k=2).indices self.assertEqual(eq_mask[idcs].sum(), 4) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test that subset_dims is called if d>=21 X_train = 1 + torch.rand(20, 21, **tkwargs) model = MockModel( MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)) ) bounds = torch.ones(2, 21, **tkwargs) bounds[1] = 2 # test NEI with X_baseline acqf = qNoisyExpectedImprovement(model, X_baseline=X_train) with mock.patch( "botorch.optim.initializers.sample_perturbed_subset_dims", wraps=sample_perturbed_subset_dims, ) as mock_subset_dims: X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=5, sigma=1e-3, bounds=bounds ) self.assertTrue(X_rnd.shape, torch.Size([5, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) mock_subset_dims.assert_called_once() # test tiny prob_perturb to make sure we perturb at least one dimension X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=5, sigma=1e-3, bounds=bounds, prob_perturb=1e-8, ) self.assertTrue( ((X_rnd.unsqueeze(0) == X_train.unsqueeze(1)).all(dim=-1)).sum() == 0 )
def test_get_X_baseline(self): tkwargs = {"device": self.device} for dtype in (torch.float, torch.double): tkwargs["dtype"] = dtype X_train = torch.rand(20, 2, **tkwargs) model = MockModel( MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True))) # test NEI with X_baseline acqf = qNoisyExpectedImprovement(model, X_baseline=X_train[:2]) X = get_X_baseline(acq_function=acqf) self.assertTrue(torch.equal(X, acqf.X_baseline)) # test EI without X_baseline acqf = qExpectedImprovement(model, best_f=0.0) with warnings.catch_warnings( record=True) as w, settings.debug(True): X_rnd = get_X_baseline(acq_function=acqf, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) # set train inputs model.train_inputs = (X_train, ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test that we fail back to train_inputs if X_baseline is an empty tensor acqf.register_buffer("X_baseline", X_train[:0]) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test acquisitipon function without X_baseline or model acqf = FixedFeatureAcquisitionFunction(acqf, d=2, columns=[0], values=[0]) with warnings.catch_warnings( record=True) as w, settings.debug(True): X_rnd = get_X_baseline(acq_function=acqf, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) Y_train = 2 * X_train[:2] + 1 moo_model = MockModel(MockPosterior(mean=Y_train, samples=Y_train)) ref_point = torch.zeros(2, **tkwargs) # test NEHVI with X_baseline acqf = qNoisyExpectedHypervolumeImprovement( moo_model, ref_point=ref_point, X_baseline=X_train[:2], cache_root=False, ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, acqf.X_baseline)) # test qEHVI without train_inputs acqf = qExpectedHypervolumeImprovement( moo_model, ref_point=ref_point, partitioning=FastNondominatedPartitioning( ref_point=ref_point, Y=Y_train, ), ) # test extracting train_inputs from model list GP model_list = ModelListGP( SingleTaskGP(X_train, Y_train[:, :1]), SingleTaskGP(X_train, Y_train[:, 1:]), ) acqf = qExpectedHypervolumeImprovement( model_list, ref_point=ref_point, partitioning=FastNondominatedPartitioning( ref_point=ref_point, Y=Y_train, ), ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test MESMO for which we need to use # `acqf.mo_model` batched_mo_model = SingleTaskGP(X_train, Y_train) acqf = qMultiObjectiveMaxValueEntropy( batched_mo_model, sample_pareto_frontiers=lambda model: torch.rand( 10, 2, **tkwargs), ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test that if there is an input transform that is applied # to the train_inputs when the model is in eval mode, we # extract the untransformed train_inputs model = SingleTaskGP(X_train, Y_train[:, :1], input_transform=Warp(indices=[0, 1])) model.eval() self.assertFalse(torch.equal(model.train_inputs[0], X_train)) acqf = qExpectedImprovement(model, best_f=0.0) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train))