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,
)