def __init__(
self,
model: Model,
sampler: Optional[MCSampler] = None,
objective: Optional[MCMultiOutputObjective] = None,
constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
X_pending: Optional[Tensor] = None,
) -> None:
r"""Constructor for the MCAcquisitionFunction base class.
Args:
model: A fitted model.
sampler: The sampler used to draw base samples. Defaults to
`SobolQMCNormalSampler(num_samples=128, collapse_batch_dims=True)`.
objective: The MCMultiOutputObjective under which the samples are
evaluated. Defaults to `IdentityMultiOutputObjective()`.
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.
X_pending: A `m x d`-dim Tensor of `m` design points that have
points that have been submitted for function evaluation
but have not yet been evaluated.
"""
super().__init__(model=model)
if sampler is None:
sampler = SobolQMCNormalSampler(num_samples=128,
collapse_batch_dims=True)
self.add_module("sampler", sampler)
if objective is None:
objective = IdentityMCMultiOutputObjective()
elif not isinstance(objective, MCMultiOutputObjective):
raise UnsupportedError(
"Only objectives of type MCMultiOutputObjective are supported for "
"Multi-Objective MC acquisition functions.")
if (hasattr(model, "input_transform")
and isinstance(model.input_transform, InputPerturbation)
and constraints is not None):
raise UnsupportedError(
"Constraints are not supported with input perturbations, due to"
"sample q-batch shape being different than that of the inputs."
"Use a composite objective that applies feasibility weighting to"
"samples before calculating the risk measure.")
self.add_module("objective", objective)
self.constraints = constraints
self.X_pending = None
if X_pending is not None:
self.set_X_pending(X_pending)
def get_botorch_objective(
model: Model,
objective_weights: Tensor,
use_scalarized_objective: bool = True,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
objective_thresholds: Optional[Tensor] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
"""Constructs a BoTorch `AcquisitionObjective` object.
Args:
model: A BoTorch Model
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
use_scalarized_objective: A boolean parameter that defaults to True,
specifying whether ScalarizedObjective should be used.
NOTE: when using outcome_constraints, use_scalarized_objective
will be ignored.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
objective_thresholds: A tensor containing thresholds forming a reference point
from which to calculate pareto frontier hypervolume. Points that do not
dominate the objective_thresholds contribute nothing to hypervolume.
X_observed: Observed points that are feasible and appear in the
objective or the constraints. None if there are no such points.
Returns:
A BoTorch `AcquisitionObjective` object. It will be one of:
`ScalarizedObjective`, `LinearMCOObjective`, `ConstrainedMCObjective`.
"""
if objective_thresholds is not None:
nonzero_idcs = torch.nonzero(objective_weights).view(-1)
objective_weights = objective_weights[nonzero_idcs]
return WeightedMCMultiOutputObjective(weights=objective_weights,
outcomes=nonzero_idcs.tolist())
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch Objective.")
if outcome_constraints:
if use_scalarized_objective:
logger.warning(
"Currently cannot use ScalarizedObjective when there are outcome "
"constraints. Ignoring (default) kwarg `use_scalarized_objective`"
"= True. Creating ConstrainedMCObjective.")
obj_tf = get_objective_weights_transform(objective_weights)
def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
return ConstrainedMCObjective(objective=objective,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
elif use_scalarized_objective:
return ScalarizedObjective(weights=objective_weights)
return LinearMCObjective(weights=objective_weights)
def posterior(self,
X: Tensor,
output_indices: Optional[List[int]] = None,
**kwargs: Any) -> DeterministicPosterior:
r"""Compute the (deterministic) posterior at X.
Args:
X: A `batch_shape x n x d`-dim input tensor `X`.
output_indices: A list of indices, corresponding to the outputs over
which to compute the posterior. If omitted, computes the posterior
over all model outputs.
Returns:
A `DeterministicPosterior` object, representing `batch_shape` joint
posteriors over `n` points and the outputs selected by `output_indices`.
"""
# Apply the input transforms in `eval` mode.
self.eval()
X = self.transform_inputs(X)
# Note: we use a Tensor instance check so that `observation_noise = True`
# just gets ignored. This avoids having to do a bunch of case distinctions
# when using a ModelList.
if isinstance(kwargs.get("observation_noise"), Tensor):
# TODO: Consider returning an MVN here instead
raise UnsupportedError(
"Deterministic models do not support observation noise.")
values = self.forward(X)
# NOTE: The `outcome_transform` `untransform`s the predictions rather than the
# `posterior` (as is done in GP models). This is more general since it works
# even if the transform doesn't support `untransform_posterior`.
if hasattr(self, "outcome_transform"):
values, _ = self.outcome_transform.untransform(values)
if output_indices is not None:
values = values[..., output_indices]
return DeterministicPosterior(values=values)
def __init__(
self,
model: Model,
sampler: Optional[MCSampler] = None,
objective: Optional[MCMultiOutputObjective] = None,
X_pending: Optional[Tensor] = None,
) -> None:
r"""Constructor for the MCAcquisitionFunction base class.
Args:
model: A fitted model.
sampler: The sampler used to draw base samples. Defaults to
`SobolQMCNormalSampler(num_samples=128, collapse_batch_dims=True)`.
objective: The MCMultiOutputObjective under which the samples are
evaluated. Defaults to `IdentityMultiOutputObjective()`.
X_pending: A `m x d`-dim Tensor of `m` design points that have
points that have been submitted for function evaluation
but have not yet been evaluated.
"""
super().__init__(model=model)
if sampler is None:
sampler = SobolQMCNormalSampler(num_samples=128,
collapse_batch_dims=True)
self.add_module("sampler", sampler)
if objective is None:
objective = IdentityMCMultiOutputObjective()
elif not isinstance(objective, MCMultiOutputObjective):
raise UnsupportedError(
"Only objectives of type MCMultiOutputObjective are supported for "
"Multi-Objective MC acquisition functions.")
self.add_module("objective", objective)
self.X_pending = None
if X_pending is not None:
self.set_X_pending(X_pending)
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]]]:
# `outcome_constraints` is validated to be None in `gen`
if outcome_constraints is not None:
raise UnsupportedError("Outcome constraints not yet supported.")
return get_out_of_sample_best_point_acqf(
model=not_none(self.model),
Xs=self.Xs,
objective_weights=objective_weights,
# With None `outcome_constraints`, `get_objective` utility
# always returns a `ScalarizedObjective`, which results in
# `get_out_of_sample_best_point_acqf` always selecting
# `PosteriorMean`.
outcome_constraints=outcome_constraints,
X_observed=not_none(X_observed),
seed_inner=seed_inner,
fixed_features=fixed_features,
fidelity_features=self.fidelity_features,
target_fidelities=target_fidelities,
qmc=qmc,
)
def get_best_f_mc(
training_data: TrainingData,
objective: Optional[MCAcquisitionObjective] = None,
posterior_transform: Optional[PosteriorTransform] = None,
) -> Tensor:
if not training_data.is_block_design:
raise NotImplementedError(
"Currently only block designs are supported.")
Y = training_data.Y
posterior_transform = _deprecate_objective_arg(
posterior_transform=posterior_transform,
objective=objective
if not isinstance(objective, MCAcquisitionObjective) else None,
)
if posterior_transform is not None:
# retain the original tensor dimension since objective expects explicit
# output dimension.
Y_dim = Y.dim()
Y = posterior_transform.evaluate(Y)
if Y.dim() < Y_dim:
Y = Y.unsqueeze(-1)
if objective is None:
if Y.shape[-1] > 1:
raise UnsupportedError(
"Acquisition functions require an objective when "
"used with multi-output models (execpt for multi-objective"
"acquisition functions).")
objective = IdentityMCObjective()
return objective(Y).max(-1).values
def test_raise_botorch_exceptions(self):
with self.assertRaises(BotorchError):
raise BotorchError("message")
with self.assertRaises(CandidateGenerationError):
raise CandidateGenerationError("message")
with self.assertRaises(UnsupportedError):
raise UnsupportedError("message")
def _deprecate_objective_arg(
posterior_transform: Optional[PosteriorTransform] = None,
objective: Optional[AcquisitionObjective] = None,
) -> Optional[PosteriorTransform]:
if posterior_transform is not None:
if objective is None:
return posterior_transform
else:
raise RuntimeError(
"Got both a non-MC objective (DEPRECATED) and a posterior "
"transform. Use only a posterior transform instead.")
elif objective is not None:
warnings.warn(
"The `objective` argument to AnalyticAcquisitionFunctions is deprecated "
"and will be removed in the next version. Use `posterior_transform` "
"instead.",
DeprecationWarning,
)
if not isinstance(objective, ScalarizedObjective):
raise UnsupportedError(
"Analytic acquisition functions only support ScalarizedObjective "
"(DEPRECATED) type objectives.")
return ScalarizedPosteriorTransform(weights=objective.weights,
offset=objective.offset)
else:
return None
def __init__(
self,
acq_function: AcquisitionFunction,
proximal_weights: Tensor,
) -> None:
r"""Derived Acquisition Function weighted by proximity to recently
observed point.
Args:
acq_function: The base acquisition function, operating on input tensors
of feature dimension `d`.
proximal_weights: A `d` dim tensor used to bias locality
along each axis.
"""
Module.__init__(self)
self.acq_func = acq_function
if hasattr(acq_function, "X_pending"):
if acq_function.X_pending is not None:
raise UnsupportedError(
"Proximal acquisition function requires `X_pending` to be None."
)
self.X_pending = acq_function.X_pending
self.register_buffer("proximal_weights", proximal_weights)
# check model for train_inputs and single batch
if not hasattr(self.acq_func.model, "train_inputs"):
raise UnsupportedError("Acquisition function model must have "
"`train_inputs`.")
if (self.acq_func.model.batch_shape != torch.Size([])
and self.acq_func.model.train_inputs[0].shape[1] != 1):
raise UnsupportedError(
"Proximal acquisition function requires a single batch model")
# check to make sure that weights match the training data shape
if (len(self.proximal_weights.shape) != 1
or self.proximal_weights.shape[0] !=
self.acq_func.model.train_inputs[0][-1].shape[-1]):
raise ValueError(
"`proximal_weights` must be a one dimensional tensor with "
"same feature dimension as model.")
def get_botorch_objective_and_transform(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
objective_thresholds: Optional[Tensor] = None,
X_observed: Optional[Tensor] = None,
) -> Tuple[Optional[MCAcquisitionObjective], Optional[PosteriorTransform]]:
"""Constructs a BoTorch `AcquisitionObjective` object.
Args:
model: A BoTorch Model
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
objective_thresholds: A tensor containing thresholds forming a reference point
from which to calculate pareto frontier hypervolume. Points that do not
dominate the objective_thresholds contribute nothing to hypervolume.
X_observed: Observed points that are feasible and appear in the
objective or the constraints. None if there are no such points.
Returns:
A two-tuple containing (optioally) an `MCAcquisitionObjective` and
(optionally) a `PosteriorTransform`.
"""
if objective_thresholds is not None:
# we are doing multi-objective optimization
nonzero_idcs = torch.nonzero(objective_weights).view(-1)
objective_weights = objective_weights[nonzero_idcs]
objective = WeightedMCMultiOutputObjective(
weights=objective_weights, outcomes=nonzero_idcs.tolist())
return objective, None
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch objective.")
if outcome_constraints:
# If there are outcome constraints, we use MC Acquistion functions
obj_tf = get_objective_weights_transform(objective_weights)
def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=objective,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return objective, None
# Case of linear weights - use ScalarizedPosteriorTransform
transform = ScalarizedPosteriorTransform(weights=objective_weights)
return None, transform
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
iteration_fidelity: Optional[int] = None,
data_fidelity: Optional[int] = None,
linear_truncated: bool = True,
nu: float = 2.5,
likelihood: Optional[Likelihood] = None,
outcome_transform: Optional[OutcomeTransform] = None,
input_transform: Optional[InputTransform] = None,
) -> None:
self._init_args = {
"iteration_fidelity": iteration_fidelity,
"data_fidelity": data_fidelity,
"linear_truncated": linear_truncated,
"nu": nu,
"outcome_transform": outcome_transform,
}
if iteration_fidelity is None and data_fidelity is None:
raise UnsupportedError(
"SingleTaskMultiFidelityGP requires at least one fidelity parameter."
)
if input_transform is not None:
input_transform.to(train_X)
with torch.no_grad():
transformed_X = self.transform_inputs(
X=train_X, input_transform=input_transform)
self._set_dimensions(train_X=transformed_X, train_Y=train_Y)
covar_module, subset_batch_dict = _setup_multifidelity_covar_module(
dim=transformed_X.size(-1),
aug_batch_shape=self._aug_batch_shape,
iteration_fidelity=iteration_fidelity,
data_fidelity=data_fidelity,
linear_truncated=linear_truncated,
nu=nu,
)
super().__init__(
train_X=train_X,
train_Y=train_Y,
likelihood=likelihood,
covar_module=covar_module,
outcome_transform=outcome_transform,
input_transform=input_transform,
)
self._subset_batch_dict = {
"likelihood.noise_covar.raw_noise": -2,
"mean_module.constant": -2,
"covar_module.raw_outputscale": -1,
**subset_batch_dict,
}
self.to(train_X)
def __init__(
self,
model: Model,
sampler: Optional[MCSampler] = None,
objective: Optional[MCAcquisitionObjective] = None,
posterior_transform: Optional[PosteriorTransform] = None,
X_pending: Optional[Tensor] = None,
) -> None:
r"""Constructor for the MCAcquisitionFunction base class.
Args:
model: A fitted model.
sampler: The sampler used to draw base samples. Defaults to
`SobolQMCNormalSampler(num_samples=512, collapse_batch_dims=True)`.
objective: The MCAcquisitionObjective under which the samples are
evaluated. Defaults to `IdentityMCObjective()`.
posterior_transform: A PosteriorTransform (optional).
X_pending: A `batch_shape, m x d`-dim Tensor of `m` design points
that have points that have been submitted for function evaluation
but have not yet been evaluated.
"""
super().__init__(model=model)
if sampler is None:
sampler = SobolQMCNormalSampler(num_samples=512,
collapse_batch_dims=True)
self.add_module("sampler", sampler)
if objective is None and model.num_outputs != 1:
if posterior_transform is None:
raise UnsupportedError(
"Must specify an objective or a posterior transform when using "
"a multi-output model.")
elif not posterior_transform.scalarize:
raise UnsupportedError(
"If using a multi-output model without an objective, "
"posterior_transform must scalarize the output.")
if objective is None:
objective = IdentityMCObjective()
self.posterior_transform = posterior_transform
self.add_module("objective", objective)
self.set_X_pending(X_pending)
def posterior(
self, X: Tensor, output_indices: Optional[List[int]] = None, **kwargs: Any
) -> DeterministicPosterior:
r"""Compute the (deterministic) posterior at X."""
if kwargs.get("observation_noise") is not None:
# TODO: Consider returning an MVN here instead
raise UnsupportedError(
"Deterministic models do not support observation noise."
)
values = self.forward(X)
if output_indices is not None:
values = values[..., output_indices]
return DeterministicPosterior(values=values)
def get_best_f_mc(
training_data: TrainingData,
objective: Optional[AcquisitionObjective] = None,
) -> Tensor:
if not training_data.is_block_design:
raise NotImplementedError("Currently only block designs are supported.")
Y = training_data.Y
if objective is None:
if Y.shape[-1] > 1:
raise UnsupportedError(
"Acquisition functions require an objective when "
"used with multi-output models (execpt for multi-objective"
"acquisition functions)."
)
objective = IdentityMCObjective()
return objective(training_data.Y).max(-1).values
def get_botorch_objective(
model: Model,
objective_weights: Tensor,
use_scalarized_objective: bool = True,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
"""Constructs a BoTorch `AcquisitionObjective` object.
Args:
model: A BoTorch Model
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
use_scalarized_objective: A boolean parameter that defaults to True,
specifying whether ScalarizedObjective should be used.
NOTE: when using outcome_constraints, use_scalarized_objective
will be ignored.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
X_observed: Observed points that are feasible and appear in the
objective or the constraints. None if there are no such points.
Returns:
A BoTorch `AcquisitionObjective` object. It will be one of:
`ScalarizedObjective`, `LinearMCOObjective`, `ConstrainedMCObjective`.
"""
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch Objective.")
if outcome_constraints:
if use_scalarized_objective:
logger.warning(
"Currently cannot use ScalarizedObjective when there are outcome "
"constraints. Ignoring (default) kwarg `use_scalarized_objective`"
"= True. Creating ConstrainedMCObjective.")
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
return ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
if use_scalarized_objective:
return ScalarizedObjective(weights=objective_weights)
return LinearMCObjective(weights=objective_weights)
def __init__(
self, model: Model, objective: Optional[AnalyticMultiOutputObjective] = None
) -> None:
r"""Constructor for the MultiObjectiveAnalyticAcquisitionFunction base class.
Args:
model: A fitted model.
objective: An AnalyticMultiOutputObjective (optional).
"""
super().__init__(model=model)
if objective is None:
objective = IdentityAnalyticMultiOutputObjective()
elif not isinstance(objective, AnalyticMultiOutputObjective):
raise UnsupportedError(
"Only objectives of type AnalyticMultiOutputObjective are supported "
"for Multi-Objective analytic acquisition functions."
)
self.objective = objective
def __init__(
self,
model: Model,
num_mv_samples: int,
posterior_transform: Optional[PosteriorTransform] = None,
maximize: bool = True,
X_pending: Optional[Tensor] = None,
) -> None:
r"""Single-outcome max-value entropy search-based acquisition functions.
Args:
model: A fitted single-outcome model.
num_mv_samples: Number of max value samples.
posterior_transform: A PosteriorTransform. If using a multi-output model,
a PosteriorTransform that transforms the multi-output posterior into a
single-output posterior is required.
maximize: If True, consider the problem a maximization problem.
X_pending: A `m x d`-dim Tensor of `m` design points that have been
submitted for function evaluation but have not yet been evaluated.
"""
super().__init__(model=model)
if posterior_transform is None and model.num_outputs != 1:
raise UnsupportedError(
"Must specify a posterior transform when using a multi-output model."
)
# Batched GP models are not currently supported
try:
batch_shape = model.batch_shape
except NotImplementedError:
batch_shape = torch.Size()
if len(batch_shape) > 0:
raise NotImplementedError(
"Batched GP models (e.g., fantasized models) are not yet "
f"supported by `{self.__class__.__name__}`."
)
self.num_mv_samples = num_mv_samples
self.posterior_transform = posterior_transform
self.maximize = maximize
self.weight = 1.0 if maximize else -1.0
self.set_X_pending(X_pending)
def _check_sampler(self) -> None:
r"""Check compatibility of sampler and model with a cached Cholesky."""
if not self.sampler.collapse_batch_dims:
raise UnsupportedError(
"Expected sampler to use `collapse_batch_dims=True`.")
elif self.sampler.base_samples is not None:
warnings.warn(
message=
("sampler.base_samples is not None. The base_samples must be "
"initialized to None. Resetting sampler.base_samples to None."
),
category=BotorchWarning,
)
self.sampler.base_samples = None
elif self._uses_matheron and self.sampler.batch_range != (0, -1):
raise RuntimeError(
"sampler.batch_range is not (0, -1). This check requires that the "
"sampler.batch_range is (0, -1) with GPs that use Matheron's rule "
"for sampling, in order to properly collapse batch dimensions. "
)
def get_botorch_objective(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
"""Constructs a BoTorch `Objective`."""
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch Objective.")
if outcome_constraints is None:
objective = ScalarizedObjective(weights=objective_weights)
else:
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return objective
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
train_Yvar: Tensor,
iteration_fidelity: Optional[int] = None,
data_fidelity: Optional[int] = None,
linear_truncated: bool = True,
nu: float = 2.5,
outcome_transform: Optional[OutcomeTransform] = None,
) -> None:
if iteration_fidelity is None and data_fidelity is None:
raise UnsupportedError(
"FixedNoiseMultiFidelityGP requires at least one fidelity parameter."
)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
covar_module, subset_batch_dict = _setup_multifidelity_covar_module(
dim=train_X.size(-1),
aug_batch_shape=self._aug_batch_shape,
iteration_fidelity=iteration_fidelity,
data_fidelity=data_fidelity,
linear_truncated=linear_truncated,
nu=nu,
)
super().__init__(
train_X=train_X,
train_Y=train_Y,
train_Yvar=train_Yvar,
covar_module=covar_module,
outcome_transform=outcome_transform,
)
self._subset_batch_dict = {
"likelihood.noise_covar.raw_noise": -2,
"mean_module.constant": -2,
"covar_module.raw_outputscale": -1,
**subset_batch_dict,
}
self.to(train_X)
def forward(self, X: Tensor) -> Tensor:
r"""Evaluate analytical EUBO on the candidate set X.
Args:
X: A `batch_shape x q x d`-dim Tensor, where `q = 2` if `previous_winner`
is not `None`, and `q = 1` otherwise.
Returns:
The acquisition value for each batch as a tensor of shape `batch_shape`.
"""
if not ((X.shape[-2] == 2) or ((X.shape[-2] == 1) and
(self.previous_winner is not None))):
raise UnsupportedError(
f"{self.__class__.__name__} only support q=2 or q=1"
"with a previous winner specified")
Y = X if self.outcome_model is None else self.outcome_model(X)
if self.previous_winner is not None:
Y = torch.cat([Y, match_batch_shape(self.previous_winner, Y)],
dim=-2)
# Calling forward directly instead of posterior here to
# obtain the full covariance matrix
pref_posterior = self.model(Y)
pref_mean = pref_posterior.mean
pref_cov = pref_posterior.covariance_matrix
delta = pref_mean[..., 0] - pref_mean[..., 1]
sigma = torch.sqrt(pref_cov[..., 0, 0] + pref_cov[..., 1, 1] -
pref_cov[..., 0, 1] - pref_cov[..., 1, 0])
u = delta / sigma
ucdf = self.std_norm.cdf(u)
updf = torch.exp(self.std_norm.log_prob(u))
acqf_val = sigma * (updf + u * ucdf)
if self.previous_winner is None:
acqf_val = acqf_val + pref_mean[..., 1]
return acqf_val
def __init__(
self,
model: Model,
num_fantasies: Optional[int] = 64,
sampler: Optional[MCSampler] = None,
objective: Optional[AcquisitionObjective] = None,
inner_sampler: Optional[MCSampler] = None,
X_pending: Optional[Tensor] = None,
current_value: Optional[Tensor] = None,
) -> None:
r"""q-Knowledge Gradient (one-shot optimization).
Args:
model: A fitted model. Must support fantasizing.
num_fantasies: The number of fantasy points to use. More fantasy
points result in a better approximation, at the expense of
memory and wall time. Unused if `sampler` is specified.
sampler: The sampler used to sample fantasy observations. Optional
if `num_fantasies` is specified.
objective: The objective under which the samples are evaluated. If
`None` or a ScalarizedObjective, then the analytic posterior mean
is used, otherwise the objective is MC-evaluated (using
inner_sampler).
inner_sampler: The sampler used for inner sampling. Ignored if the
objective is `None` or a ScalarizedObjective.
X_pending: A `m x d`-dim Tensor of `m` design points that have
points that have been submitted for function evaluation
but have not yet been evaluated.
current_value: The current value, i.e. the expected best objective
given the observed points `D`. If omitted, forward will not
return the actual KG value, but the expected best objective
given the data set `D u X`.
"""
if sampler is None:
if num_fantasies is None:
raise ValueError(
"Must specify `num_fantasies` if no `sampler` is provided."
)
# base samples should be fixed for joint optimization over X, X_fantasies
sampler = SobolQMCNormalSampler(
num_samples=num_fantasies, resample=False, collapse_batch_dims=True
)
elif num_fantasies is not None:
if sampler.sample_shape != torch.Size([num_fantasies]):
raise ValueError(
f"The sampler shape must match num_fantasies={num_fantasies}."
)
else:
num_fantasies = sampler.sample_shape[0]
super(MCAcquisitionFunction, self).__init__(model=model)
# if not explicitly specified, we use the posterior mean for linear objs
if isinstance(objective, MCAcquisitionObjective) and inner_sampler is None:
inner_sampler = SobolQMCNormalSampler(
num_samples=128, resample=False, collapse_batch_dims=True
)
if objective is None and model.num_outputs != 1:
raise UnsupportedError(
"Must specify an objective when using a multi-output model."
)
self.sampler = sampler
self.objective = objective
self.set_X_pending(X_pending)
self.inner_sampler = inner_sampler
self.num_fantasies = num_fantasies
self.current_value = current_value
def __init__(
self,
model: Model,
num_fantasies: Optional[int] = 64,
sampler: Optional[MCSampler] = None,
objective: Optional[AcquisitionObjective] = None,
inner_sampler: Optional[MCSampler] = None,
X_pending: Optional[Tensor] = None,
current_value: Optional[Tensor] = None,
cost_aware_utility: Optional[CostAwareUtility] = None,
project: Callable[[Tensor], Tensor] = lambda X: X,
expand: Callable[[Tensor], Tensor] = lambda X: X,
) -> None:
r"""Multi-Fidelity q-Knowledge Gradient (one-shot optimization).
Args:
model: A fitted model. Must support fantasizing.
num_fantasies: The number of fantasy points to use. More fantasy
points result in a better approximation, at the expense of
memory and wall time. Unused if `sampler` is specified.
sampler: The sampler used to sample fantasy observations. Optional
if `num_fantasies` is specified.
objective: The objective under which the samples are evaluated. If
`None` or a ScalarizedObjective, then the analytic posterior mean
is used, otherwise the objective is MC-evaluated (using
inner_sampler).
inner_sampler: The sampler used for inner sampling. Ignored if the
objective is `None` or a ScalarizedObjective.
X_pending: A `m x d`-dim Tensor of `m` design points that have
points that have been submitted for function evaluation
but have not yet been evaluated.
current_value: The current value, i.e. the expected best objective
given the observed points `D`. If omitted, forward will not
return the actual KG value, but the expected best objective
given the data set `D u X`.
cost_aware_utility: A CostAwareUtility computing the cost-transformed
utility from a candidate set and samples of increases in utility.
project: A callable mapping a `batch_shape x q x d` tensor of design
points to a tensor of the same shape projected to the desired
target set (e.g. the target fidelities in case of multi-fidelity
optimization).
expand: A callable mapping a `batch_shape x q x d` input tensor to
a `batch_shape x (q + q_e)' x d`-dim output tensor, where the
`q_e` additional points in each q-batch correspond to
additional ("trace") observations.
"""
if current_value is None and cost_aware_utility is not None:
raise UnsupportedError(
"Cost-aware KG requires current_value to be specified."
)
super().__init__(
model=model,
num_fantasies=num_fantasies,
sampler=sampler,
objective=objective,
inner_sampler=inner_sampler,
X_pending=X_pending,
current_value=current_value,
)
self.cost_aware_utility = cost_aware_utility
self.project = project
self.expand = expand
self._cost_sampler = None
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 (predicted by the given surrogate model) point
and instantiates it.
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(Keys.SUBSET_MODEL, True):
subset_model_results = subset_model(
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
)
model = subset_model_results.model
objective_weights = subset_model_results.objective_weights
outcome_constraints = subset_model_results.outcome_constraints
fixed_features = fixed_features or {}
target_fidelities = target_fidelities or {}
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."
)
acqf_class, acqf_options = pick_best_out_of_sample_point_acqf_class(
outcome_constraints=outcome_constraints,
mc_samples=mc_samples,
qmc=qmc,
seed_inner=seed_inner,
)
objective, posterior_transform = get_botorch_objective_and_transform(
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
)
if objective is not None:
if not isinstance(objective, MCAcquisitionObjective):
raise UnsupportedError(
f"Unknown objective type: {objective.__class__}" # pragma: nocover
)
acqf_options = {"objective": objective, **acqf_options}
if posterior_transform is not None:
acqf_options = {
"posterior_transform": posterior_transform,
**acqf_options
}
acqf = acqf_class(model=model, **acqf_options) # pyre-ignore [45]
if fixed_features:
acqf = FixedFeatureAcquisitionFunction(
acq_function=acqf,
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 acqf, non_fixed_idcs
def gen(
self,
n: int,
bounds: List,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata,
Optional[List[TCandidateMetadata]]]:
r"""Generate new candidates.
Args:
n: Number of candidates to generate.
bounds: A list of (lower, upper) tuples for each column of X.
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b.
linear_constraints: A tuple of (A, b). For k linear constraints on
d-dimensional x, A is (k x d) and b is (k x 1) such that
A x <= b.
fixed_features: A map {feature_index: value} for features that
should be fixed to a particular value during generation.
pending_observations: A list of m (k_i x d) feature tensors X
for m outcomes and k_i pending observations for outcome i.
model_gen_options: A config dictionary that can contain
model-specific options.
rounding_func: A function that rounds an optimization result
appropriately (i.e., according to `round-trip` transformations).
target_fidelities: A map {feature_index: value} of fidelity feature
column indices to their respective target fidelities. Used for
multi-fidelity optimization.
Returns:
3-element tuple containing
- (n x d) tensor of generated points.
- n-tensor of weights for each point.
- Dictionary of model-specific metadata for the given
generation candidates.
"""
options = model_gen_options or {}
acf_options = options.get("acquisition_function_kwargs", {})
optimizer_options = options.get("optimizer_kwargs", {})
X_pending, X_observed = _get_X_pending_and_observed(
Xs=self.Xs,
pending_observations=pending_observations,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
bounds=bounds,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
)
# subset model only to the outcomes we need for the optimization
model = not_none(self.model)
if options.get("subset_model", True):
model, objective_weights, outcome_constraints, _ = subset_model(
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
)
objective = get_botorch_objective(
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
)
inequality_constraints = _to_inequality_constraints(linear_constraints)
# TODO: update optimizers to handle inequality_constraints
if inequality_constraints is not None:
raise UnsupportedError(
"Inequality constraints are not yet supported for KnowledgeGradient!"
)
# extract a few options
n_fantasies = acf_options.get("num_fantasies", 64)
qmc = acf_options.get("qmc", True)
seed_inner = acf_options.get("seed_inner", None)
num_restarts = optimizer_options.get("num_restarts", 40)
raw_samples = optimizer_options.get("raw_samples", 1024)
# get current value
current_value = self._get_current_value(
model=model,
bounds=bounds,
X_observed=not_none(X_observed),
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
seed_inner=seed_inner,
fixed_features=fixed_features,
model_gen_options=model_gen_options,
target_fidelities=target_fidelities,
qmc=qmc,
)
bounds_ = torch.tensor(bounds, dtype=self.dtype, device=self.device)
bounds_ = bounds_.transpose(0, 1)
# get acquisition function
acq_function = _instantiate_KG(
model=model,
objective=objective,
qmc=qmc,
n_fantasies=n_fantasies,
num_trace_observations=options.get("num_trace_observations", 0),
mc_samples=acf_options.get("mc_samples", 256),
seed_inner=seed_inner,
seed_outer=acf_options.get("seed_outer", None),
X_pending=X_pending,
target_fidelities=target_fidelities,
fidelity_weights=options.get("fidelity_weights"),
current_value=current_value,
cost_intercept=self.cost_intercept,
)
# optimize and get new points
new_x = _optimize_and_get_candidates(
acq_function=acq_function,
bounds_=bounds_,
n=n,
num_restarts=num_restarts,
raw_samples=raw_samples,
optimizer_options=optimizer_options,
rounding_func=rounding_func,
inequality_constraints=inequality_constraints,
fixed_features=fixed_features,
)
return new_x, torch.ones(n, dtype=self.dtype), {}, None
def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
raise UnsupportedError(
"Analytic acquisition functions do not account for X_pending yet."
)
def prune_inferior_points(
model: Model,
X: Tensor,
objective: Optional[MCAcquisitionObjective] = None,
num_samples: int = 2048,
max_frac: float = 1.0,
) -> Tensor:
r"""Prune points from an input tensor that are unlikely to be the best point.
Given a model, an objective, and an input tensor `X`, this function returns
the subset of points in `X` that have some probability of being the best
point under the objective. This function uses sampling to estimate the
probabilities, the higher the number of points `n` in `X` the higher the
number of samples `num_samples` should be to obtain accurate estimates.
Args:
model: A fitted model. Batched models are currently not supported.
X: An input tensor of shape `n x d`. Batched inputs are currently not
supported.
objective: The objective under which to evaluate the posterior.
num_samples: The number of samples used to compute empirical
probabilities of being the best point.
max_frac: The maximum fraction of points to retain. Must satisfy
`0 < max_frac <= 1`. Ensures that the number of elements in the
returned tensor does not exceed `ceil(max_frac * n)`.
Returns:
A `n' x d` with subset of points in `X`, where
n' = min(N_nz, ceil(max_frac * n))
with `N_nz` the number of points in `X` that have non-zero (empirical,
under `num_samples` samples) probability of being the best point.
"""
if X.ndim > 2:
# TODO: support batched inputs (req. dealing with ragged tensors)
raise UnsupportedError(
"Batched inputs `X` are currently unsupported by prune_inferior_points"
)
max_points = math.ceil(max_frac * X.size(-2))
if max_points < 1 or max_points > X.size(-2):
raise ValueError(f"max_frac must take values in (0, 1], is {max_frac}")
with torch.no_grad():
posterior = model.posterior(X=X)
if posterior.event_shape.numel() > SobolEngine.MAXDIM:
if settings.debug.on():
warnings.warn(
f"Sample dimension q*m={posterior.event_shape.numel()} exceeding Sobol "
f"max dimension ({SobolEngine.MAXDIM}). Using iid samples instead.",
SamplingWarning,
)
sampler = IIDNormalSampler(num_samples=num_samples)
else:
sampler = SobolQMCNormalSampler(num_samples=num_samples)
samples = sampler(posterior)
if objective is None:
objective = IdentityMCObjective()
obj_vals = objective(samples)
if obj_vals.ndim > 2:
# TODO: support batched inputs (req. dealing with ragged tensors)
raise UnsupportedError(
"Batched models are currently unsupported by prune_inferior_points"
)
is_best = torch.argmax(obj_vals, dim=-1)
idcs, counts = torch.unique(is_best, return_counts=True)
if len(idcs) > max_points:
counts, order_idcs = torch.sort(counts, descending=True)
idcs = order_idcs[:max_points]
return X[idcs]
def _make_linear_constraints(
indices: Tensor,
coefficients: Tensor,
rhs: float,
shapeX: torch.Size,
eq: bool = False,
) -> List[ScipyConstraintDict]:
r"""Create linear constraints to be used by `scipy.minimize`.
Encodes constraints of the form
`\sum_i (coefficients[i] * X[..., indices[i]]) ? rhs`
where `?` can be designated either as `>=` by setting `eq=False`, or as
`=` by setting `eq=True`.
If indices is one-dimensional, the constraints are broadcasted across
all elements of the q-batch. If indices is two-dimensional, then
constraints are applied across elements of a q-batch. In either case,
constraints are created for all t-batches.
Args:
indices: A tensor of shape `c` or `c x 2`, where c is the number of terms
in the constraint. If single-dimensional, contains the indices of
the dimensions of the feature space that occur in the linear
constraint. If two-dimensional, contains pairs of indices of the
q-batch (0) and the feature space (1) that occur in the linear
constraint.
coefficients: A single-dimensional tensor of coefficients with the same
number of elements as `indices`.
rhs: The right hand side of the constraint.
shapeX: The shape of the torch tensor to construct the constraints for
(i.e. `b x q x d`). Must have three dimensions.
eq: If True, return an equality constraint, o/w return an inequality
constraint (indicated by "eq" / "ineq" value of the `type` key).
Returns:
A list of constraint dictionaries with the following keys
- "type": Indicates the type of the constraint ("eq" if `eq=True`, "ineq" o/w)
- "fun": A callable evaluating the constraint value on `x`, a flattened
version of the input tensor `X`, returning a scalar.
- "jac": A callable evaluating the constraint's Jacobian on `x`, a flattened
version of the input tensor `X`, returning a numpy array.
"""
if len(shapeX) != 3:
raise UnsupportedError("`shapeX` must be `b x q x d`")
q, d = shapeX[-2:]
n = shapeX.numel()
constraints: List[ScipyConstraintDict] = []
coeffs = _arrayify(coefficients)
ctype = "eq" if eq else "ineq"
if indices.dim() > 2:
raise UnsupportedError(
"Linear constraints supported only on individual candidates and "
"across q-batches, not across general batch shapes.")
elif indices.dim() == 2:
# indices has two dimensions (potential constraints across q-batch elements)
if indices[:, 0].max() > q - 1:
raise RuntimeError(f"Index out of bounds for {q}-batch")
if indices[:, 1].max() > d - 1:
raise RuntimeError(
f"Index out of bounds for {d}-dim parameter tensor")
offsets = [shapeX[i:].numel() for i in range(1, len(shapeX))]
# rule is [i, j, k] is at
# i * offsets[0] + j * offsets[1] + k
for i in range(shapeX[0]):
idxr = []
for a in indices:
b = a.tolist()
idxr.append(i * offsets[0] + b[0] * offsets[1] + b[1])
fun = partial(eval_lin_constraint,
flat_idxr=idxr,
coeffs=coeffs,
rhs=float(rhs))
jac = partial(lin_constraint_jac,
flat_idxr=idxr,
coeffs=coeffs,
n=n)
constraints.append({"type": ctype, "fun": fun, "jac": jac})
elif indices.dim() == 1:
# indices is one-dim - broadcast constraints across q-batches and t-batches
if indices.max() > d - 1:
raise RuntimeError(
f"Index out of bounds for {d}-dim parameter tensor")
offsets = [shapeX[i:].numel() for i in range(1, len(shapeX))]
for i in range(shapeX[0]):
for j in range(shapeX[1]):
idxr = (i * offsets[0] + j * offsets[1] + indices).tolist()
fun = partial(eval_lin_constraint,
flat_idxr=idxr,
coeffs=coeffs,
rhs=float(rhs))
jac = partial(lin_constraint_jac,
flat_idxr=idxr,
coeffs=coeffs,
n=n)
constraints.append({"type": ctype, "fun": fun, "jac": jac})
else:
raise ValueError("`indices` must be at least one-dimensional")
return constraints
def _pad_batch_pareto_frontier(
Y: Tensor,
ref_point: Tensor,
is_pareto: bool = False,
feasibility_mask: Optional[Tensor] = None,
) -> Tensor:
r"""Get a batch Pareto frontier by padding the pareto frontier with repeated points.
This assumes maximization.
Args:
Y: A `(batch_shape) x n x m`-dim tensor of points
ref_point: a `(batch_shape) x m`-dim tensor containing the reference point
is_pareto: a boolean indicating whether the points in Y are already
non-dominated.
feasibility_mask: A `(batch_shape) x n`-dim tensor of booleans indicating
whether each point is feasible.
Returns:
A `(batch_shape) x max_num_pareto x m`-dim tensor of padded Pareto
frontiers.
"""
tkwargs = {"dtype": Y.dtype, "device": Y.device}
ref_point = ref_point.unsqueeze(-2)
batch_shape = Y.shape[:-2]
if len(batch_shape) > 1:
raise UnsupportedError(
"_pad_batch_pareto_frontier only supports a single "
f"batch dimension, but got {len(batch_shape)} "
"batch dimensions.")
if feasibility_mask is not None:
# set infeasible points to be the reference point (corresponding to the batch)
Y = torch.where(feasibility_mask.unsqueeze(-1), Y, ref_point)
if not is_pareto:
pareto_mask = is_non_dominated(Y)
else:
pareto_mask = torch.ones(Y.shape[:-1],
dtype=torch.bool,
device=Y.device)
better_than_ref = (Y > ref_point).all(dim=-1)
# is_non_dominated assumes maximization
# TODO: filter out points that are worse than the reference point first here
pareto_mask = pareto_mask & better_than_ref
if len(batch_shape) == 0:
return Y[pareto_mask]
# Note: in the batch case, the Pareto frontier is padded by repeating
# a Pareto point. This ensures that the padded box-decomposition has
# the same number of points, which enables fast batch operations.
max_n_pareto = pareto_mask.sum(dim=-1).max().item()
pareto_Y = torch.empty(*batch_shape, max_n_pareto, Y.shape[-1], **tkwargs)
for i, pareto_i in enumerate(pareto_mask):
pareto_i = Y[i, pareto_mask[i]]
n_pareto = pareto_i.shape[0]
if n_pareto > 0:
pareto_Y[i, :n_pareto] = pareto_i
# pad pareto_Y, so that all batches have the same size Pareto set
pareto_Y[i, n_pareto:] = pareto_i[-1]
else:
# if there are no pareto points in this batch, use the reference
# point
pareto_Y[i, :] = ref_point[i]
return pareto_Y
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
cat_dims: List[int],
cont_kernel_factory: Optional[Callable[[int, List[int]], Kernel]] = None,
likelihood: Optional[Likelihood] = None,
outcome_transform: Optional[OutcomeTransform] = None, # TODO
input_transform: Optional[InputTransform] = None, # TODO
) -> None:
r"""A single-task exact GP model supporting categorical parameters.
Args:
train_X: A `batch_shape x n x d` tensor of training features.
train_Y: A `batch_shape x n x m` tensor of training observations.
cat_dims: A list of indices corresponding to the columns of
the input `X` that should be considered categorical features.
cont_kernel_factory: A method that accepts `ard_num_dims` and
`active_dims` arguments and returns an instatiated GPyTorch
`Kernel` object to be used as the ase kernel for the continuous
dimensions. If omitted, this model uses a Matern-2.5 kernel as
the kernel for the ordinal parameters.
likelihood: A likelihood. If omitted, use a standard
GaussianLikelihood with inferred noise level.
# outcome_transform: An outcome transform that is applied to the
# training data during instantiation and to the posterior during
# inference (that is, the `Posterior` obtained by calling
# `.posterior` on the model will be on the original scale).
# input_transform: An input transform that is applied in the model's
# forward pass.
Example:
>>> train_X = torch.cat(
[torch.rand(20, 2), torch.randint(3, (20, 1))], dim=-1)
)
>>> train_Y = (
torch.sin(train_X[..., :-1]).sum(dim=1, keepdim=True)
+ train_X[..., -1:]
)
>>> model = MixedSingleTaskGP(train_X, train_Y, cat_dims=[-1])
"""
if outcome_transform is not None:
raise UnsupportedError("outcome transforms not yet supported")
if input_transform is not None:
raise UnsupportedError("input transforms not yet supported")
if len(cat_dims) == 0:
raise ValueError(
"Must specify categorical dimensions for MixedSingleTaskGP"
)
input_batch_shape, aug_batch_shape = self.get_batch_dimensions(
train_X=train_X, train_Y=train_Y
)
if cont_kernel_factory is None:
def cont_kernel_factory(
batch_shape: torch.Size, ard_num_dims: int, active_dims: List[int]
) -> MaternKernel:
return MaternKernel(
nu=2.5,
batch_shape=batch_shape,
ard_num_dims=ard_num_dims,
active_dims=active_dims,
)
if likelihood is None:
# This Gamma prior is quite close to the Horseshoe prior
min_noise = 1e-5 if train_X.dtype == torch.float else 1e-6
likelihood = GaussianLikelihood(
batch_shape=aug_batch_shape,
noise_constraint=GreaterThan(
min_noise, transform=None, initial_value=1e-3
),
noise_prior=GammaPrior(0.9, 10.0),
)
d = train_X.shape[-1]
cat_dims = normalize_indices(indices=cat_dims, d=d)
ord_dims = sorted(set(range(d)) - set(cat_dims))
if len(ord_dims) == 0:
covar_module = ScaleKernel(
CategoricalKernel(
batch_shape=aug_batch_shape,
ard_num_dims=len(cat_dims),
)
)
else:
sum_kernel = ScaleKernel(
cont_kernel_factory(
batch_shape=aug_batch_shape,
ard_num_dims=len(ord_dims),
active_dims=ord_dims,
)
+ ScaleKernel(
CategoricalKernel(
batch_shape=aug_batch_shape,
ard_num_dims=len(cat_dims),
active_dims=cat_dims,
)
)
)
prod_kernel = ScaleKernel(
cont_kernel_factory(
batch_shape=aug_batch_shape,
ard_num_dims=len(ord_dims),
active_dims=ord_dims,
)
* CategoricalKernel(
batch_shape=aug_batch_shape,
ard_num_dims=len(cat_dims),
active_dims=cat_dims,
)
)
covar_module = sum_kernel + prod_kernel
super().__init__(
train_X=train_X,
train_Y=train_Y,
likelihood=likelihood,
covar_module=covar_module,
outcome_transform=outcome_transform,
input_transform=input_transform,
)
