def _compute_qehvi(self,
samples: Tensor,
X: Optional[Tensor] = None) -> Tensor:
r"""Compute the expected (feasible) hypervolume improvement given MC samples.
Args:
samples: A `n_samples x batch_shape x q' x m`-dim tensor of samples.
X: A `batch_shape x q x d`-dim tensor of inputs.
Returns:
A `batch_shape x (model_batch_shape)`-dim tensor of expected hypervolume
improvement for each batch.
"""
# Note that the objective may subset the outcomes (e.g. this will usually happen
# if there are constraints present).
obj = self.objective(samples, X=X)
q = obj.shape[-2]
if self.constraints is not None:
feas_weights = torch.ones(obj.shape[:-1],
device=obj.device,
dtype=obj.dtype)
feas_weights = apply_constraints_nonnegative_soft(
obj=feas_weights,
constraints=self.constraints,
samples=samples,
eta=self.eta,
)
self._cache_q_subset_indices(q_out=q)
batch_shape = obj.shape[:-2]
# this is n_samples x input_batch_shape x
areas_per_segment = torch.zeros(
*batch_shape,
self.cell_lower_bounds.shape[-2],
dtype=obj.dtype,
device=obj.device,
)
cell_batch_ndim = self.cell_lower_bounds.ndim - 2
sample_batch_view_shape = torch.Size([
batch_shape[0] if cell_batch_ndim > 0 else 1,
*[1 for _ in range(len(batch_shape) - max(cell_batch_ndim, 1))],
*self.cell_lower_bounds.shape[1:-2],
])
view_shape = (
*sample_batch_view_shape,
self.cell_upper_bounds.shape[-2],
1,
self.cell_upper_bounds.shape[-1],
)
for i in range(1, self.q_out + 1):
# TODO: we could use batches to compute (q choose i) and (q choose q-i)
# simultaneously since subsets of size i and q-i have the same number of
# elements. This would decrease the number of iterations, but increase
# memory usage.
q_choose_i = self.q_subset_indices[f"q_choose_{i}"]
# this tensor is mc_samples x batch_shape x i x q_choose_i x m
obj_subsets = obj.index_select(dim=-2, index=q_choose_i.view(-1))
obj_subsets = obj_subsets.view(obj.shape[:-2] + q_choose_i.shape +
obj.shape[-1:])
# since all hyperrectangles share one vertex, the opposite vertex of the
# overlap is given by the component-wise minimum.
# take the minimum in each subset
overlap_vertices = obj_subsets.min(dim=-2).values
# add batch-dim to compute area for each segment (pseudo-pareto-vertex)
# this tensor is mc_samples x batch_shape x num_cells x q_choose_i x m
overlap_vertices = torch.min(
overlap_vertices.unsqueeze(-3),
self.cell_upper_bounds.view(view_shape))
# substract cell lower bounds, clamp min at zero
lengths_i = (
overlap_vertices -
self.cell_lower_bounds.view(view_shape)).clamp_min(0.0)
# take product over hyperrectangle side lengths to compute area
# sum over all subsets of size i
areas_i = lengths_i.prod(dim=-1)
# if constraints are present, apply a differentiable approximation of
# the indicator function
if self.constraints is not None:
feas_subsets = feas_weights.index_select(
dim=-1,
index=q_choose_i.view(-1)).view(feas_weights.shape[:-1] +
q_choose_i.shape)
areas_i = areas_i * feas_subsets.unsqueeze(-3).prod(dim=-1)
areas_i = areas_i.sum(dim=-1)
# Using the inclusion-exclusion principle, set the sign to be positive
# for subsets of odd sizes and negative for subsets of even size
areas_per_segment += (-1)**(i + 1) * areas_i
# sum over segments and average over MC samples
return areas_per_segment.sum(dim=-1).mean(dim=0)
def _compute_qehvi(self, samples: Tensor) -> Tensor:
r"""Compute the expected (feasible) hypervolume improvement given MC samples.
Args:
samples: A `n_samples x batch_shape x q x m`-dim tensor of samples.
Returns:
A `batch_shape`-dim tensor of expected hypervolume improvement for each
batch.
"""
q = samples.shape[-2]
# Note that the objective may subset the outcomes (e.g. this will usually happen
# if there are constraints present).
obj = self.objective(samples)[..., : self.ref_point.shape[0]]
if self.constraints is not None:
feas_weights = torch.ones(
obj.shape[:-1], device=obj.device, dtype=obj.dtype
)
feas_weights = apply_constraints_nonnegative_soft(
obj=feas_weights,
constraints=self.constraints,
samples=samples,
eta=self.eta,
)
self._cache_q_subset_indices(q=q)
batch_shape = samples.shape[:-2]
areas_per_segment = torch.zeros(
*batch_shape,
self.cell_lower_bounds.shape[-2],
dtype=obj.dtype,
device=obj.device,
)
sample_batch_view_shape = [
batch_shape[0] if self.cell_lower_bounds.ndim == 3 else 1
] + [1] * (len(batch_shape) - 1)
view_shape = (
*sample_batch_view_shape,
self.cell_upper_bounds.shape[-2],
1,
self.cell_upper_bounds.shape[-1],
)
for i in range(1, q + 1):
# TODO: we could use batches to compute (q choose i) and (q choose q-i)
# simulataneously since subsets of size i and q-i have the same number of
# elements. This would decrease the number of iterations, but increase
# memory usage.
q_choose_i = self.q_subset_indices[f"q_choose_{i}"]
# this tensor is mc_samples x batch_shape x i x q_choose_i x m
obj_subsets = torch.stack(
[obj.index_select(dim=-2, index=q_choose_i[:, k]) for k in range(i)],
dim=-3,
)
# since all hyperrectangles share one vertex, the opposite vertex of the
# overlap is given by the component-wise minimum.
# take the minimum in each subset
overlap_vertices = obj_subsets.min(dim=-3).values
expanded_shape = (
batch_shape
+ self.cell_upper_bounds.shape[-2:-1]
+ overlap_vertices.shape[-2:]
)
# add batch-dim to compute area for each segment (pseudo-pareto-vertex)
# this tensor is mc_samples x batch_shape x num_cells x q_choose_i x m
overlap_vertices = overlap_vertices.unsqueeze(-3).expand(
*batch_shape,
self.cell_lower_bounds.shape[-2],
*overlap_vertices.shape[-2:],
)
overlap_vertices = torch.min(
overlap_vertices,
self.cell_upper_bounds.view(view_shape).expand(expanded_shape),
)
# substract cell lower bounds, clamp min at zero
lengths_i = overlap_vertices - self.cell_lower_bounds.view(
view_shape
).expand(expanded_shape)
lengths_i = lengths_i.clamp_min(0.0)
# take product over hyperrectangle side lengths to compute area
# sum over all subsets of size i
areas_i = lengths_i.prod(dim=-1).sum(dim=-1)
# if constraints are present, apply a differentiable approximation of
# the indicator function
if self.constraints is not None:
feas_subsets = torch.cat(
[
feas_weights.index_select(dim=-1, index=q_choose_i[:, k])
for k in range(i)
],
dim=-1,
)
areas_i = areas_i * feas_subsets.prod(-1, keepdim=True)
# Using the inclusion-exclusion principle, set the sign to be positive
# for subsets of odd sizes and negative for subsets of even size
areas_per_segment += (-1) ** (i + 1) * areas_i
# sum over segments and average over MC samples
return areas_per_segment.sum(dim=-1).mean(dim=0)