def test_q_upper_confidence_bound(self): for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 1 x 1 x 1 samples = torch.zeros(1, 1, 1, device=self.device, dtype=dtype) mm = MockModel(MockPosterior(samples=samples)) # X is `q x d` = 1 x 1. X is a dummy and unused b/c of mocking X = torch.zeros(1, 1, device=self.device, dtype=dtype) # basic test sampler = IIDNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) # basic test, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1])) bs = acqf.sampler.base_samples.clone() res = acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) # basic test for X_pending and warning acqf.set_X_pending() self.assertIsNone(acqf.X_pending) acqf.set_X_pending(None) self.assertIsNone(acqf.X_pending) acqf.set_X_pending(X) self.assertEqual(acqf.X_pending, X) res = acqf(X) X2 = torch.zeros( 1, 1, 1, device=self.device, dtype=dtype, requires_grad=True ) with warnings.catch_warnings(record=True) as ws, settings.debug(True): acqf.set_X_pending(X2) self.assertEqual(acqf.X_pending, X2) self.assertEqual(len(ws), 1) self.assertTrue(issubclass(ws[-1].category, BotorchWarning))
def test_q_upper_confidence_bound(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 1 x 1 x 1 samples = torch.zeros(1, 1, 1, device=device, dtype=dtype) mm = MockModel(MockPosterior(samples=samples)) # X is `q x d` = 1 x 1. X is a dummy and unused b/c of mocking X = torch.zeros(1, 1, device=device, dtype=dtype) # basic test sampler = IIDNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) # basic test, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1])) bs = acqf.sampler.base_samples.clone() res = acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def test_acquisition_functions(self): tkwargs = {"device": self.device, "dtype": torch.double} train_X, train_Y, train_Yvar, model = self._get_data_and_model( infer_noise=True, **tkwargs ) fit_fully_bayesian_model_nuts( model, warmup_steps=8, num_samples=5, thinning=2, disable_progbar=True ) sampler = IIDNormalSampler(num_samples=2) acquisition_functions = [ ExpectedImprovement(model=model, best_f=train_Y.max()), ProbabilityOfImprovement(model=model, best_f=train_Y.max()), PosteriorMean(model=model), UpperConfidenceBound(model=model, beta=4), qExpectedImprovement(model=model, best_f=train_Y.max(), sampler=sampler), qNoisyExpectedImprovement(model=model, X_baseline=train_X, sampler=sampler), qProbabilityOfImprovement( model=model, best_f=train_Y.max(), sampler=sampler ), qSimpleRegret(model=model, sampler=sampler), qUpperConfidenceBound(model=model, beta=4, sampler=sampler), qNoisyExpectedHypervolumeImprovement( model=ModelListGP(model, model), X_baseline=train_X, ref_point=torch.zeros(2, **tkwargs), sampler=sampler, ), qExpectedHypervolumeImprovement( model=ModelListGP(model, model), ref_point=torch.zeros(2, **tkwargs), sampler=sampler, partitioning=NondominatedPartitioning( ref_point=torch.zeros(2, **tkwargs), Y=train_Y.repeat([1, 2]) ), ), ] for acqf in acquisition_functions: for batch_shape in [[5], [6, 5, 2]]: test_X = torch.rand(*batch_shape, 1, 4, **tkwargs) self.assertEqual(acqf(test_X).shape, torch.Size(batch_shape))
def get_acquisition_function( acquisition_function_name: str, model: Model, objective: MCAcquisitionObjective, X_observed: Tensor, X_pending: Optional[Tensor] = None, mc_samples: int = 500, qmc: bool = True, seed: Optional[int] = None, **kwargs, ) -> monte_carlo.MCAcquisitionFunction: r"""Convenience function for initializing botorch acquisition functions. Args: acquisition_function_name: Name of the acquisition function. model: A fitted model. objective: A MCAcquisitionObjective. X_observed: A `m1 x d`-dim Tensor of `m1` design points that have already been observed. X_pending: A `m2 x d`-dim Tensor of `m2` design points whose evaluation is pending. mc_samples: The number of samples to use for (q)MC evaluation of the acquisition function. qmc: If True, use quasi-Monte-Carlo sampling (instead of iid). seed: If provided, perform deterministic optimization (i.e. the function to optimize is fixed and not stochastic). Returns: The requested acquisition function. Example: >>> model = SingleTaskGP(train_X, train_Y) >>> obj = LinearMCObjective(weights=torch.tensor([1.0, 2.0])) >>> acqf = get_acquisition_function("qEI", model, obj, train_X) """ # initialize the sampler if qmc: sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed) else: sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed) # instantiate and return the requested acquisition function if acquisition_function_name == "qEI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qExpectedImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, ) elif acquisition_function_name == "qPI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qProbabilityOfImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, tau=kwargs.get("tau", 1e-3), ) elif acquisition_function_name == "qNEI": return monte_carlo.qNoisyExpectedImprovement( model=model, X_baseline=X_observed, sampler=sampler, objective=objective, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", False), ) elif acquisition_function_name == "qSR": return monte_carlo.qSimpleRegret(model=model, sampler=sampler, objective=objective, X_pending=X_pending) elif acquisition_function_name == "qUCB": if "beta" not in kwargs: raise ValueError("`beta` must be specified in kwargs for qUCB.") return monte_carlo.qUpperConfidenceBound( model=model, beta=kwargs["beta"], sampler=sampler, objective=objective, X_pending=X_pending, ) raise NotImplementedError( f"Unknown acquisition function {acquisition_function_name}")
def get_acquisition_function( acquisition_function_name: str, model: Model, objective: MCAcquisitionObjective, X_observed: Tensor, X_pending: Optional[Tensor] = None, constraints: Optional[List[Callable[[Tensor], Tensor]]] = None, mc_samples: int = 500, qmc: bool = True, seed: Optional[int] = None, **kwargs, ) -> monte_carlo.MCAcquisitionFunction: r"""Convenience function for initializing botorch acquisition functions. Args: acquisition_function_name: Name of the acquisition function. model: A fitted model. objective: A MCAcquisitionObjective. X_observed: A `m1 x d`-dim Tensor of `m1` design points that have already been observed. X_pending: A `m2 x d`-dim Tensor of `m2` design points whose evaluation is pending. constraints: A list of callables, each mapping a Tensor of dimension `sample_shape x batch-shape x q x m` to a Tensor of dimension `sample_shape x batch-shape x q`, where negative values imply feasibility. Used when constraint_transforms are not passed as part of the objective. mc_samples: The number of samples to use for (q)MC evaluation of the acquisition function. qmc: If True, use quasi-Monte-Carlo sampling (instead of iid). seed: If provided, perform deterministic optimization (i.e. the function to optimize is fixed and not stochastic). Returns: The requested acquisition function. Example: >>> model = SingleTaskGP(train_X, train_Y) >>> obj = LinearMCObjective(weights=torch.tensor([1.0, 2.0])) >>> acqf = get_acquisition_function("qEI", model, obj, train_X) """ # initialize the sampler if qmc: sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed) else: sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed) # instantiate and return the requested acquisition function if acquisition_function_name == "qEI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qExpectedImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, ) elif acquisition_function_name == "qPI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qProbabilityOfImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, tau=kwargs.get("tau", 1e-3), ) elif acquisition_function_name == "qNEI": return monte_carlo.qNoisyExpectedImprovement( model=model, X_baseline=X_observed, sampler=sampler, objective=objective, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", False), ) elif acquisition_function_name == "qSR": return monte_carlo.qSimpleRegret(model=model, sampler=sampler, objective=objective, X_pending=X_pending) elif acquisition_function_name == "qUCB": if "beta" not in kwargs: raise ValueError("`beta` must be specified in kwargs for qUCB.") return monte_carlo.qUpperConfidenceBound( model=model, beta=kwargs["beta"], sampler=sampler, objective=objective, X_pending=X_pending, ) elif acquisition_function_name == "qEHVI": # pyre-fixme [16]: `Model` has no attribute `train_targets` try: ref_point = kwargs["ref_point"] except KeyError: raise ValueError( "`ref_point` must be specified in kwargs for qEHVI") try: Y = kwargs["Y"] except KeyError: raise ValueError("`Y` must be specified in kwargs for qEHVI") # get feasible points if constraints is not None: feas = torch.stack([c(Y) <= 0 for c in constraints], dim=-1).all(dim=-1) Y = Y[feas] obj = objective(Y) partitioning = NondominatedPartitioning( ref_point=torch.as_tensor(ref_point, dtype=Y.dtype, device=Y.device), Y=obj, alpha=kwargs.get("alpha", 0.0), ) return moo_monte_carlo.qExpectedHypervolumeImprovement( model=model, ref_point=ref_point, partitioning=partitioning, sampler=sampler, objective=objective, constraints=constraints, X_pending=X_pending, ) raise NotImplementedError( f"Unknown acquisition function {acquisition_function_name}")
def test_q_upper_confidence_bound_batch(self): # TODO: T41739913 Implement tests for all MCAcquisitionFunctions for dtype in (torch.float, torch.double): samples = torch.zeros(2, 2, 1, device=self.device, dtype=dtype) samples[0, 0, 0] = 1.0 mm = MockModel(MockPosterior(samples=samples)) # X is a dummy and unused b/c of mocking X = torch.zeros(1, 1, 1, device=self.device, dtype=dtype) # test batch mode sampler = IIDNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # test batch mode, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test batch mode, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test batch mode, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) # basic test for X_pending and warning acqf.set_X_pending() self.assertIsNone(acqf.X_pending) acqf.set_X_pending(None) self.assertIsNone(acqf.X_pending) acqf.set_X_pending(X) self.assertEqual(acqf.X_pending, X) res = acqf(X) X2 = torch.zeros( 1, 1, 1, device=self.device, dtype=dtype, requires_grad=True ) with warnings.catch_warnings(record=True) as ws, settings.debug(True): acqf.set_X_pending(X2) self.assertEqual(acqf.X_pending, X2) self.assertEqual(len(ws), 1) self.assertTrue(issubclass(ws[-1].category, BotorchWarning))
# model.likelihood.noise_covar.register_constraint("raw_noise", GreaterThan(1e1)) for it in range(25): t0 = time.time() # Negate the first element of the objective, since botorch maximizes. objective = ScalarizedObjective(torch.tensor([-1.0] + [0.0] * dim_basis)) # q: number of candidates. q = 10 # candidate_x has shape (q, dim_basis) t1 = time.time() candidate_x, acq_value = optimize_acqf( # UpperConfidenceBound(model, beta=0.1, objective=objective), qUpperConfidenceBound(model, beta=0.1, objective=objective), bounds=torch.Tensor([[-1] * dim_basis, [1] * dim_basis]), q=q, # Not sure why these are necessary: num_restarts=1, raw_samples=1, ) print(f"optimize_acqf took {time.time()-t1} secs") # Evaluate candidate point. candidate_y = jax_to_torch([loss_and_grad(torch_to_jax(candidate_x[i, :])) for i in range(q)]) train_x = torch.cat([train_x, candidate_x]) train_y = torch.cat([train_y, candidate_y]) # This is currently an error "Cannot yet add fantasy observations to multitask GPs, but this is coming soon!" # model = model.condition_on_observations(X=candidate_x, Y=candidate_y)
def test_q_upper_confidence_bound_batch(self, cuda=False): # TODO: T41739913 Implement tests for all MCAcquisitionFunctions device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): samples = torch.zeros(2, 2, 1, device=device, dtype=dtype) samples[0, 0, 0] = 1.0 mm = MockModel(MockPosterior(samples=samples)) # X is a dummy and unused b/c of mocking X = torch.zeros(1, 1, 1, device=device, dtype=dtype) # test batch mode sampler = IIDNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # test batch mode, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test batch mode, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test batch mode, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True) acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))