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)