def test_posterior_mean(self): for dtype in (torch.float, torch.double): mean = torch.tensor([[0.25]], device=self.device, dtype=dtype) mm = MockModel(MockPosterior(mean=mean)) module = PosteriorMean(model=mm) X = torch.empty(1, 1, device=self.device, dtype=dtype) pm = module(X) self.assertTrue(torch.equal(pm, mean.view(-1))) # check for proper error if multi-output model mean2 = torch.rand(1, 2, device=self.device, dtype=dtype) mm2 = MockModel(MockPosterior(mean=mean2)) with self.assertRaises(UnsupportedError): PosteriorMean(model=mm2)
def test_posterior_mean_batch(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): mean = torch.tensor([-0.5, 0.0, 0.5], device=device, dtype=dtype).view(3, 1, 1) mm = MockModel(MockPosterior(mean=mean)) module = PosteriorMean(model=mm) X = torch.empty(3, 1, 1, device=device, dtype=dtype) pm = module(X) self.assertTrue(torch.equal(pm, mean.view(-1))) # check for proper error if multi-output model mean2 = torch.rand(3, 1, 2, device=device, dtype=dtype) mm2 = MockModel(MockPosterior(mean=mean2)) module2 = PosteriorMean(model=mm2) with self.assertRaises(UnsupportedError): module2(X)
def _get_value_function( model: Model, objective: Optional[Union[MCAcquisitionObjective, ScalarizedObjective]] = None, sampler: Optional[MCSampler] = None, project: Optional[Callable[[Tensor], Tensor]] = None, valfunc_cls: Optional[Type[AcquisitionFunction]] = None, valfunc_argfac: Optional[Callable[[Model, Dict[str, Any]]]] = None, ) -> AcquisitionFunction: r"""Construct value function (i.e. inner acquisition function).""" if valfunc_cls is not None: common_kwargs: Dict[str, Any] = {"model": model, "objective": objective} if issubclass(valfunc_cls, MCAcquisitionFunction): common_kwargs["sampler"] = sampler kwargs = valfunc_argfac(model=model) if valfunc_argfac is not None else {} base_value_function = valfunc_cls(**common_kwargs, **kwargs) else: if isinstance(objective, MCAcquisitionObjective): base_value_function = qSimpleRegret( model=model, sampler=sampler, objective=objective ) else: base_value_function = PosteriorMean(model=model, objective=objective) if project is None: return base_value_function else: return ProjectedAcquisitionFunction( base_value_function=base_value_function, project=project, )
def _get_best_point_acqf( self, X_observed: Tensor, objective_weights: Tensor, mc_samples: int = 512, fixed_features: Optional[Dict[int, float]] = None, target_fidelities: Optional[Dict[int, float]] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, seed_inner: Optional[int] = None, qmc: bool = True, **kwargs: Any, ) -> Tuple[AcquisitionFunction, Optional[List[int]]]: model = self.model # subset model only to the outcomes we need for the optimization if kwargs.get("subset_model", True): model, objective_weights, outcome_constraints = subset_model( model=model, # pyre-ignore: [6] objective_weights=objective_weights, outcome_constraints=outcome_constraints, ) fixed_features = fixed_features or {} target_fidelities = target_fidelities or {} objective = ScalarizedObjective(weights=objective_weights) acq_function = PosteriorMean( model=model, objective=objective # pyre-ignore: [6] ) if self.fidelity_features: # we need to optimize at the target fidelities if any(f in self.fidelity_features for f in fixed_features): raise RuntimeError("Fixed features cannot also be fidelity features") elif not set(self.fidelity_features) == set(target_fidelities): raise RuntimeError( "Must provide a target fidelity for every fidelity feature" ) # make sure to not modify fixed_features in-place fixed_features = {**fixed_features, **target_fidelities} elif target_fidelities: raise RuntimeError( "Must specify fidelity_features in fit() when using target fidelities" ) if fixed_features: acq_function = FixedFeatureAcquisitionFunction( acq_function=acq_function, d=X_observed.size(-1), columns=list(fixed_features.keys()), values=list(fixed_features.values()), ) non_fixed_idcs = [ i for i in range(self.Xs[0].size(-1)) if i not in fixed_features ] else: non_fixed_idcs = None return acq_function, non_fixed_idcs
def _get_value_function( model: Model, objective: Optional[Union[MCAcquisitionObjective, ScalarizedObjective]] = None, sampler: Optional[MCSampler] = None, ) -> AcquisitionFunction: r"""Construct value function (i.e. inner acquisition function).""" if isinstance(objective, MCAcquisitionObjective): return qSimpleRegret(model=model, sampler=sampler, objective=objective) else: return PosteriorMean(model=model, objective=objective)
def test_init(self): NO = "botorch.utils.testing.MockModel.num_outputs" with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 1 mm = MockModel(None) acqf = PosteriorMean(mm) BS = BoltzmannSampling(acqf) self.assertEqual(BS.acq_func, acqf) self.assertEqual(BS.eta, 1.0) self.assertTrue(BS.replacement) BS = BoltzmannSampling(acqf, eta=0.5, replacement=False) self.assertEqual(BS.acq_func, acqf) self.assertEqual(BS.eta, 0.5) self.assertFalse(BS.replacement)
def _get_best_point_acqf( self, X_observed: Tensor, objective_weights: Tensor, fixed_features: Optional[Dict[int, float]] = None, target_fidelities: Optional[Dict[int, float]] = None, ) -> Tuple[AcquisitionFunction, Optional[List[int]]]: fixed_features = fixed_features or {} target_fidelities = target_fidelities or {} objective = ScalarizedObjective(weights=objective_weights) acq_function = PosteriorMean( model=self.model, objective=objective # pyre-ignore: [6] ) if self.fidelity_features: # we need to optimize at the target fidelities if any(f in self.fidelity_features for f in fixed_features): raise RuntimeError("Fixed features cannot also be fidelity features") elif not set(self.fidelity_features) == set(target_fidelities): raise RuntimeError( "Must provide a target fidelity for every fidelity feature" ) # make sure to not modify fixed_features in-place fixed_features = {**fixed_features, **target_fidelities} elif target_fidelities: raise RuntimeError( "Must specify fidelity_features in fit() when using target fidelities" ) if fixed_features: acq_function = FixedFeatureAcquisitionFunction( acq_function=acq_function, d=X_observed.size(-1), columns=list(fixed_features.keys()), values=list(fixed_features.values()), ) non_fixed_idcs = [ i for i in range(self.Xs[0].size(-1)) if i not in fixed_features ] else: non_fixed_idcs = None return acq_function, non_fixed_idcs
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_value_function( model: Model, objective: Optional[Union[MCAcquisitionObjective, ScalarizedObjective]] = None, sampler: Optional[MCSampler] = None, project: Optional[Callable[[Tensor], Tensor]] = None, ) -> AcquisitionFunction: r"""Construct value function (i.e. inner acquisition function).""" if isinstance(objective, MCAcquisitionObjective): base_value_function = qSimpleRegret(model=model, sampler=sampler, objective=objective) else: base_value_function = PosteriorMean(model=model, objective=objective) if project is None: return base_value_function else: return ProjectedAcquisitionFunction( base_value_function=base_value_function, project=project, )
def test_KnowledgeGradient(self): model = KnowledgeGradient() model.fit( Xs=self.Xs, Ys=self.Ys, Yvars=self.Yvars, search_space_digest=SearchSpaceDigest( feature_names=self.feature_names, bounds=self.bounds, ), metric_names=self.metric_names, ) n = 2 X_dummy = torch.rand(1, n, 4, dtype=self.dtype, device=self.device) acq_dummy = torch.tensor(0.0, dtype=self.dtype, device=self.device) with mock.patch(self.optimize_acqf) as mock_optimize_acqf: mock_optimize_acqf.side_effect = [(X_dummy, acq_dummy)] Xgen, wgen, _, __ = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=None, model_gen_options={ "acquisition_function_kwargs": self.acq_options, "optimizer_kwargs": self.optimizer_options, }, ) self.assertTrue(torch.equal(Xgen, X_dummy.cpu())) self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=self.dtype))) # called once, the best point call is not caught by mock mock_optimize_acqf.assert_called_once() ini_dummy = torch.rand(10, 32, 3, dtype=self.dtype, device=self.device) optimizer_options2 = { "num_restarts": 1, "raw_samples": 1, "maxiter": 5, "batch_limit": 1, "partial_restarts": 2, } with mock.patch( "ax.models.torch.botorch_kg.gen_one_shot_kg_initial_conditions", return_value=ini_dummy, ) as mock_warmstart_initialization: Xgen, wgen, _, __ = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=None, model_gen_options={ "acquisition_function_kwargs": self.acq_options, "optimizer_kwargs": optimizer_options2, }, ) mock_warmstart_initialization.assert_called_once() obj = ScalarizedObjective(weights=self.objective_weights) dummy_acq = PosteriorMean(model=model.model, objective=obj) with mock.patch("ax.models.torch.utils.PosteriorMean", return_value=dummy_acq) as mock_posterior_mean: Xgen, wgen, _, __ = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=None, model_gen_options={ "acquisition_function_kwargs": self.acq_options, "optimizer_kwargs": optimizer_options2, }, ) self.assertEqual(mock_posterior_mean.call_count, 2) # Check best point selection within bounds (some numerical tolerance) xbest = model.best_point(bounds=self.bounds, objective_weights=self.objective_weights) lb = torch.tensor([b[0] for b in self.bounds]) - 1e-5 ub = torch.tensor([b[1] for b in self.bounds]) + 1e-5 self.assertTrue(torch.all(xbest <= ub)) self.assertTrue(torch.all(xbest >= lb)) # test error message linear_constraints = ( torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]), torch.tensor([[0.5], [1.0]]), ) with self.assertRaises(UnsupportedError): Xgen, wgen = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=linear_constraints, ) # test input warping self.assertFalse(model.use_input_warping) model = KnowledgeGradient(use_input_warping=True) model.fit( Xs=self.Xs, Ys=self.Ys, Yvars=self.Yvars, search_space_digest=SearchSpaceDigest( feature_names=self.feature_names, bounds=self.bounds, ), metric_names=self.metric_names, ) self.assertTrue(model.use_input_warping) self.assertTrue(hasattr(model.model, "input_transform")) self.assertIsInstance(model.model.input_transform, Warp) # test loocv pseudo likelihood self.assertFalse(model.use_loocv_pseudo_likelihood) model = KnowledgeGradient(use_loocv_pseudo_likelihood=True) model.fit( Xs=self.Xs, Ys=self.Ys, Yvars=self.Yvars, search_space_digest=SearchSpaceDigest( feature_names=self.feature_names, bounds=self.bounds, ), metric_names=self.metric_names, ) self.assertTrue(model.use_loocv_pseudo_likelihood)
def _get_best_point_acqf( self, X_observed: Tensor, objective_weights: Tensor, mc_samples: int = 512, fixed_features: Optional[Dict[int, float]] = None, target_fidelities: Optional[Dict[int, float]] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, seed_inner: Optional[int] = None, qmc: bool = True, ) -> Tuple[AcquisitionFunction, Optional[List[int]]]: model = self.model fixed_features = fixed_features or {} target_fidelities = target_fidelities or {} objective = _get_objective( model=model, # pyre-ignore [6] objective_weights=objective_weights, outcome_constraints=outcome_constraints, X_observed=X_observed, ) if isinstance(objective, ScalarizedObjective): acq_function = PosteriorMean( model=model, objective=objective # pyre-ignore: [6] ) elif isinstance(objective, MCAcquisitionObjective): if qmc: sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed_inner) else: sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed_inner) acq_function = qSimpleRegret( model=model, sampler=sampler, objective=objective # pyre-ignore [6] ) else: raise UnsupportedError( f"Unknown objective type: {objective.__class__}" # pragma: nocover ) if self.fidelity_features: # we need to optimize at the target fidelities if any(f in self.fidelity_features for f in fixed_features): raise RuntimeError( "Fixed features cannot also be fidelity features") elif not set(self.fidelity_features) == set(target_fidelities): raise RuntimeError( "Must provide a target fidelity for every fidelity feature" ) # make sure to not modify fixed_features in-place fixed_features = {**fixed_features, **target_fidelities} elif target_fidelities: raise RuntimeError( "Must specify fidelity_features in fit() when using target fidelities" ) if fixed_features: acq_function = FixedFeatureAcquisitionFunction( acq_function=acq_function, d=X_observed.size(-1), columns=list(fixed_features.keys()), values=list(fixed_features.values()), ) non_fixed_idcs = [ i for i in range(self.Xs[0].size(-1)) if i not in fixed_features ] else: non_fixed_idcs = None return acq_function, non_fixed_idcs
def get_out_of_sample_best_point_acqf( model: Model, Xs: List[Tensor], X_observed: Tensor, objective_weights: Tensor, mc_samples: int = 512, fixed_features: Optional[Dict[int, float]] = None, fidelity_features: Optional[List[int]] = None, target_fidelities: Optional[Dict[int, float]] = None, outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, seed_inner: Optional[int] = None, qmc: bool = True, **kwargs: Any, ) -> Tuple[AcquisitionFunction, Optional[List[int]]]: """Picks an appropriate acquisition function to find the best out-of-sample (not observed) point using the given surrogate model. NOTE: Typically the appropriate function is the posterior mean, but can differ to account for fidelities etc. """ model = model # subset model only to the outcomes we need for the optimization if kwargs.get("subset_model", True): model, objective_weights, outcome_constraints = subset_model( model=model, objective_weights=objective_weights, outcome_constraints=outcome_constraints, ) fixed_features = fixed_features or {} target_fidelities = target_fidelities or {} objective = get_botorch_objective( model=model, objective_weights=objective_weights, outcome_constraints=outcome_constraints, X_observed=X_observed, ) if isinstance(objective, ScalarizedObjective): acq_function = PosteriorMean(model=model, objective=objective) elif isinstance(objective, MCAcquisitionObjective): if qmc: sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed_inner) else: sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed_inner) acq_function = qSimpleRegret(model=model, sampler=sampler, objective=objective) else: raise UnsupportedError( f"Unknown objective type: {objective.__class__}" # pragma: nocover ) if fidelity_features: # we need to optimize at the target fidelities if any(f in fidelity_features for f in fixed_features): raise RuntimeError( "Fixed features cannot also be fidelity features") elif set(fidelity_features) != set(target_fidelities): raise RuntimeError( "Must provide a target fidelity for every fidelity feature") # make sure to not modify fixed_features in-place fixed_features = {**fixed_features, **target_fidelities} elif target_fidelities: raise RuntimeError( "Must specify fidelity_features in fit() when using target fidelities" ) if fixed_features: acq_function = FixedFeatureAcquisitionFunction( acq_function=acq_function, d=X_observed.size(-1), columns=list(fixed_features.keys()), values=list(fixed_features.values()), ) non_fixed_idcs = [ i for i in range(Xs[0].size(-1)) if i not in fixed_features ] else: non_fixed_idcs = None return acq_function, non_fixed_idcs
def test_gen_value_function_initial_conditions(self): num_fantasies = 2 num_solutions = 3 num_restarts = 4 raw_samples = 5 n_train = 6 dim = 2 dtype = torch.float # run a thorough test with dtype float train_X = torch.rand(n_train, dim, device=self.device, dtype=dtype) train_Y = torch.rand(n_train, 1, device=self.device, dtype=dtype) model = SingleTaskGP(train_X, train_Y) fant_X = torch.rand(num_solutions, 1, dim, device=self.device, dtype=dtype) fantasy_model = model.fantasize(fant_X, IIDNormalSampler(num_fantasies)) bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype) value_function = PosteriorMean(fantasy_model) # test option error with self.assertRaises(ValueError): gen_value_function_initial_conditions( acq_function=value_function, bounds=bounds, num_restarts=num_restarts, raw_samples=raw_samples, current_model=model, options={"frac_random": 2.0}, ) # test output shape ics = gen_value_function_initial_conditions( acq_function=value_function, bounds=bounds, num_restarts=num_restarts, raw_samples=raw_samples, current_model=model, ) self.assertEqual( ics.shape, torch.Size([num_restarts, num_fantasies, num_solutions, 1, dim])) # test bounds self.assertTrue(torch.all(ics >= bounds[0])) self.assertTrue(torch.all(ics <= bounds[1])) # test dtype self.assertEqual(dtype, ics.dtype) # minimal test cases for when all raw samples are random, with dtype double dtype = torch.double n_train = 2 dim = 1 num_solutions = 1 train_X = torch.rand(n_train, dim, device=self.device, dtype=dtype) train_Y = torch.rand(n_train, 1, device=self.device, dtype=dtype) model = SingleTaskGP(train_X, train_Y) fant_X = torch.rand(1, 1, dim, device=self.device, dtype=dtype) fantasy_model = model.fantasize(fant_X, IIDNormalSampler(num_fantasies)) bounds = torch.tensor([[0], [1]], device=self.device, dtype=dtype) value_function = PosteriorMean(fantasy_model) ics = gen_value_function_initial_conditions( acq_function=value_function, bounds=bounds, num_restarts=1, raw_samples=1, current_model=model, options={"frac_random": 0.99}, ) self.assertEqual(ics.shape, torch.Size([1, num_fantasies, num_solutions, 1, dim])) # test bounds self.assertTrue(torch.all(ics >= bounds[0])) self.assertTrue(torch.all(ics <= bounds[1])) # test dtype self.assertEqual(dtype, ics.dtype)
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 optimize_objective( model: Model, bounds: Tensor, q: int, objective: Optional[MCAcquisitionObjective] = None, posterior_transform: Optional[PosteriorTransform] = None, linear_constraints: Optional[Tuple[Tensor, Tensor]] = None, fixed_features: Optional[Dict[int, float]] = None, target_fidelities: Optional[Dict[int, float]] = None, qmc: bool = True, mc_samples: int = 512, seed_inner: Optional[int] = None, optimizer_options: Dict[str, Any] = None, post_processing_func: Optional[Callable[[Tensor], Tensor]] = None, batch_initial_conditions: Optional[Tensor] = None, sequential: bool = False, **ignore, ) -> Tuple[Tensor, Tensor]: r"""Optimize an objective under the given model. Args: model: The model to be used in the objective. bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`. q: The cardinality of input sets on which the objective is to be evaluated. objective: The objective to optimize. posterior_transform: The posterior transform to be used in the acquisition function. linear_constraints: A tuple of (A, b). Given `k` linear constraints on a `d`-dimensional space, `A` is `k x d` and `b` is `k x 1` such that `A x <= b`. (Not used by single task models). fixed_features: A dictionary of feature assignments `{feature_index: value}` to hold fixed during generation. target_fidelities: A dictionary mapping input feature indices to fidelity values. Defaults to `{-1: 1.0}`. qmc: Toggle for enabling (qmc=1) or disabling (qmc=0) use of Quasi Monte Carlo. mc_samples: Integer number of samples used to estimate Monte Carlo objectives. seed_inner: Integer seed used to initialize the sampler passed to MCObjective. optimizer_options: Table used to lookup keyword arguments for the optimizer. post_processing_func: A function that post-processes an optimization result appropriately (i.e. according to `round-trip` transformations). batch_initial_conditions: A Tensor of initial values for the optimizer. sequential: If False, uses joint optimization, otherwise uses sequential optimization. Returns: A tuple containing the best input locations and corresponding objective values. """ if optimizer_options is None: optimizer_options = {} if objective is not None: sampler_cls = SobolQMCNormalSampler if qmc else IIDNormalSampler acq_function = qSimpleRegret( model=model, objective=objective, posterior_transform=posterior_transform, sampler=sampler_cls(num_samples=mc_samples, seed=seed_inner), ) else: acq_function = PosteriorMean(model=model, posterior_transform=posterior_transform) if fixed_features: acq_function = FixedFeatureAcquisitionFunction( acq_function=acq_function, d=bounds.shape[-1], columns=list(fixed_features.keys()), values=list(fixed_features.values()), ) free_feature_dims = list(range(len(bounds)) - fixed_features.keys()) free_feature_bounds = bounds[:, free_feature_dims] # (2, d' <= d) else: free_feature_bounds = bounds if linear_constraints is None: inequality_constraints = None else: A, b = linear_constraints inequality_constraints = [] k, d = A.shape for i in range(k): indicies = A[i, :].nonzero(as_tuple=False).squeeze() coefficients = -A[i, indicies] rhs = -b[i, 0] inequality_constraints.append((indicies, coefficients, rhs)) return optimize_acqf( acq_function=acq_function, bounds=free_feature_bounds, q=q, num_restarts=optimizer_options.get("num_restarts", 60), raw_samples=optimizer_options.get("raw_samples", 1024), options={ "batch_limit": optimizer_options.get("batch_limit", 8), "maxiter": optimizer_options.get("maxiter", 200), "nonnegative": optimizer_options.get("nonnegative", False), "method": optimizer_options.get("method", "L-BFGS-B"), }, inequality_constraints=inequality_constraints, fixed_features=None, # handled inside the acquisition function post_processing_func=post_processing_func, batch_initial_conditions=batch_initial_conditions, return_best_only=True, sequential=sequential, )
def test_KnowledgeGradient(self): model = KnowledgeGradient() model.fit( Xs=self.Xs, Ys=self.Ys, Yvars=self.Yvars, bounds=self.bounds, feature_names=self.feature_names, metric_names=self.metric_names, task_features=[], fidelity_features=[], ) n = 2 best_point_dummy = torch.rand(1, 3, dtype=self.dtype, device=self.device) X_dummy = torch.rand(1, n, 4, dtype=self.dtype, device=self.device) acq_dummy = torch.tensor(0.0, dtype=self.dtype, device=self.device) with mock.patch(self.optimize_acqf) as mock_optimize_acqf: mock_optimize_acqf.side_effect = [ (best_point_dummy, None), (X_dummy, acq_dummy), ] Xgen, wgen, _ = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=None, model_gen_options={ "acquisition_function_kwargs": self.acq_options, "optimizer_kwargs": self.optimizer_options, }, ) self.assertTrue(torch.equal(Xgen, X_dummy.cpu())) self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=self.dtype))) mock_optimize_acqf.assert_called( ) # called twice, once for best_point ini_dummy = torch.rand(10, 32, 3, dtype=self.dtype, device=self.device) optimizer_options2 = { "num_restarts": 1, "raw_samples": 1, "maxiter": 5, "batch_limit": 1, "partial_restarts": 2, } with mock.patch( "ax.models.torch.botorch_kg.gen_one_shot_kg_initial_conditions", return_value=ini_dummy, ) as mock_warmstart_initialization: Xgen, wgen, _ = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=None, model_gen_options={ "acquisition_function_kwargs": self.acq_options, "optimizer_kwargs": optimizer_options2, }, ) mock_warmstart_initialization.assert_called_once() obj = ScalarizedObjective(weights=self.objective_weights) dummy_acq = PosteriorMean(model=model.model, objective=obj) with mock.patch("ax.models.torch.botorch_kg.PosteriorMean", return_value=dummy_acq) as mock_posterior_mean: Xgen, wgen, _ = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=None, model_gen_options={ "acquisition_function_kwargs": self.acq_options, "optimizer_kwargs": optimizer_options2, }, ) self.assertEqual(mock_posterior_mean.call_count, 2) # Check best point selection X_dummy = torch.rand(3) acq_dummy = torch.tensor(0.0) with mock.patch(self.optimize_acqf, return_value=(X_dummy, acq_dummy)) as mock_optimize_acqf: xbest = model.best_point(bounds=self.bounds, objective_weights=self.objective_weights) self.assertTrue(torch.equal(xbest, X_dummy)) mock_optimize_acqf.assert_called_once() # test error message linear_constraints = ( torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]), torch.tensor([[0.5], [1.0]]), ) with self.assertRaises(UnsupportedError): Xgen, wgen = model.gen( n=n, bounds=self.bounds, objective_weights=self.objective_weights, outcome_constraints=None, linear_constraints=linear_constraints, )