def test_logger(self): # Verify log statements are properly captured # assertLogs() captures all log calls, ignoring the severity level with self.assertLogs(logger="botorch", level="INFO") as logs_cm: logger.info("Hello World!") logger.error("Goodbye Universe!") self.assertEqual( logs_cm.output, ["INFO:botorch:Hello World!", "ERROR:botorch:Goodbye Universe!"], )
def optimize_acqf( acq_function: AcquisitionFunction, bounds: Tensor, q: int, num_restarts: int, raw_samples: int, options: Optional[Dict[str, Union[bool, float, int, str]]] = None, inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, fixed_features: Optional[Dict[int, float]] = None, post_processing_func: Optional[Callable[[Tensor], Tensor]] = None, batch_initial_conditions: Optional[Tensor] = None, return_best_only: bool = True, sequential: bool = False, ) -> Tuple[Tensor, Tensor]: r"""Generate a set of candidates via multi-start optimization. Args: acq_function: An AcquisitionFunction. bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`. q: The number of candidates. num_restarts: The number of starting points for multistart acquisition function optimization. raw_samples: The number of samples for initialization. options: Options for candidate generation. inequality constraints: A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form `\sum_i (X[indices[i]] * coefficients[i]) >= rhs` equality constraints: A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form `\sum_i (X[indices[i]] * coefficients[i]) = rhs` fixed_features: A map `{feature_index: value}` for features that should be fixed to a particular value during generation. post_processing_func: A function that post-processes an optimization result appropriately (i.e., according to `round-trip` transformations). batch_initial_conditions: A tensor to specify the initial conditions. Set this if you do not want to use default initialization strategy. return_best_only: If False, outputs the solutions corresponding to all random restart initializations of the optimization. sequential: If False, uses joint optimization, otherwise uses sequential optimization. Returns: A two-element tuple containing - a `(num_restarts) x q x d`-dim tensor of generated candidates. - a tensor of associated acquisiton values. If `sequential=False`, this is a `(num_restarts)`-dim tensor of joint acquisition values (with explicit restart dimension if `return_best_only=False`). If `sequential=True`, this is a `q`-dim tensor of expected acquisition values conditional on having observed canidates `0,1,...,i-1`. Example: >>> # generate `q=2` candidates jointly using 20 random restarts >>> # and 512 raw samples >>> candidates, acq_value = optimize_acqf(qEI, bounds, 2, 20, 512) >>> generate `q=3` candidates sequentially using 15 random restarts >>> # and 256 raw samples >>> qEI = qExpectedImprovement(model, best_f=0.2) >>> bounds = torch.tensor([[0.], [1.]]) >>> candidates, acq_value_list = optimize_acqf( >>> qEI, bounds, 3, 15, 256, sequential=True >>> ) """ if sequential and q > 1: if not return_best_only: raise NotImplementedError( "return_best_only=False only supported for joint optimization" ) if isinstance(acq_function, OneShotAcquisitionFunction): raise NotImplementedError( "sequential optimization currently not supported for one-shot " "acquisition functions. Must have `sequential=False`." ) candidate_list, acq_value_list = [], [] candidates = torch.tensor([]) base_X_pending = acq_function.X_pending for i in range(q): candidate, acq_value = optimize_acqf( acq_function=acq_function, bounds=bounds, q=1, num_restarts=num_restarts, raw_samples=raw_samples, options=options or {}, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, fixed_features=fixed_features, post_processing_func=post_processing_func, batch_initial_conditions=None, return_best_only=True, sequential=False, ) candidate_list.append(candidate) acq_value_list.append(acq_value) candidates = torch.cat(candidate_list, dim=-2) acq_function.set_X_pending( torch.cat([base_X_pending, candidates], dim=-2) if base_X_pending is not None else candidates ) logger.info(f"Generated sequential candidate {i+1} of {q}") # Reset acq_func to previous X_pending state acq_function.set_X_pending(base_X_pending) return candidates, torch.stack(acq_value_list) options = options or {} if batch_initial_conditions is None: ic_gen = ( gen_one_shot_kg_initial_conditions if isinstance(acq_function, qKnowledgeGradient) else gen_batch_initial_conditions ) # TODO: Generating initial candidates should use parameter constraints. batch_initial_conditions = ic_gen( acq_function=acq_function, bounds=bounds, q=q, num_restarts=num_restarts, raw_samples=raw_samples, options=options, ) batch_limit: int = options.get("batch_limit", num_restarts) batch_candidates_list: List[Tensor] = [] batch_acq_values_list: List[Tensor] = [] start_idcs = list(range(0, num_restarts, batch_limit)) for start_idx in start_idcs: end_idx = min(start_idx + batch_limit, num_restarts) # optimize using random restart optimization batch_candidates_curr, batch_acq_values_curr = gen_candidates_scipy( initial_conditions=batch_initial_conditions[start_idx:end_idx], acquisition_function=acq_function, lower_bounds=bounds[0], upper_bounds=bounds[1], options={ k: v for k, v in options.items() if k not in ("batch_limit", "nonnegative") }, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, fixed_features=fixed_features, ) batch_candidates_list.append(batch_candidates_curr) batch_acq_values_list.append(batch_acq_values_curr) logger.info(f"Generated candidate batch {start_idx+1} of {len(start_idcs)}.") batch_candidates = torch.cat(batch_candidates_list) batch_acq_values = torch.cat(batch_acq_values_list) if post_processing_func is not None: batch_candidates = post_processing_func(batch_candidates) if return_best_only: best = torch.argmax(batch_acq_values.view(-1), dim=0) batch_candidates = batch_candidates[best] batch_acq_values = batch_acq_values[best] if isinstance(acq_function, OneShotAcquisitionFunction): batch_candidates = acq_function.extract_candidates(X_full=batch_candidates) return batch_candidates, batch_acq_values
def gen( self, num_points: int, # Current implementation only generates 1 point at a time model: MonotonicRejectionGP, ): """Query next point(s) to run by optimizing the acquisition function. Args: num_points (int, optional): Number of points to query. model (AEPsychMixin): Fitted model of the data. Returns: np.ndarray: Next set of point(s) to evaluate, [num_points x dim]. """ options = self.model_gen_options or {} num_restarts = options.get("num_restarts", 10) raw_samples = options.get("raw_samples", 1000) verbosity_freq = options.get("verbosity_freq", -1) lr = options.get("lr", 0.01) momentum = options.get("momentum", 0.9) nesterov = options.get("nesterov", True) epochs = options.get("epochs", 50) milestones = options.get("milestones", [25, 40]) gamma = options.get("gamma", 0.1) loss_constraint_fun = options.get( "loss_constraint_fun", default_loss_constraint_fun ) # Augment bounds with deriv indicator bounds = torch.cat((model.bounds_, torch.zeros(2, 1)), dim=1) # Fix deriv indicator to 0 during optimization fixed_features = {(bounds.shape[1] - 1): 0.0} # Fix explore features to random values if self.explore_features is not None: for idx in self.explore_features: val = ( bounds[0, idx] + torch.rand(1, dtype=bounds.dtype) * (bounds[1, idx] - bounds[0, idx]) ).item() fixed_features[idx] = val bounds[0, idx] = val bounds[1, idx] = val acqf = self._instantiate_acquisition_fn(model) # Initialize batch_initial_conditions = gen_batch_initial_conditions( acq_function=acqf, bounds=bounds, q=1, num_restarts=num_restarts, raw_samples=raw_samples, ) clamped_candidates = columnwise_clamp( X=batch_initial_conditions, lower=bounds[0], upper=bounds[1] ).requires_grad_(True) candidates = fix_features(clamped_candidates, fixed_features) optimizer = torch.optim.SGD( params=[clamped_candidates], lr=lr, momentum=momentum, nesterov=nesterov ) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=milestones, gamma=gamma ) # Optimize for epoch in range(epochs): loss = -acqf(candidates).sum() # adjust loss based on constraints on candidates loss = loss_constraint_fun(loss, candidates) if verbosity_freq > 0 and epoch % verbosity_freq == 0: logger.info("Iter: {} - Value: {:.3f}".format(epoch, -(loss.item()))) def closure(): optimizer.zero_grad() loss.backward( retain_graph=True ) # Variational model requires retain_graph return loss optimizer.step(closure) clamped_candidates.data = columnwise_clamp( X=clamped_candidates, lower=bounds[0], upper=bounds[1] ) candidates = fix_features(clamped_candidates, fixed_features) lr_scheduler.step() # Extract best point with torch.no_grad(): batch_acquisition = acqf(candidates) best = torch.argmax(batch_acquisition.view(-1), dim=0) Xopt = candidates[best][:, :-1].detach() return Xopt
def optimize_acqf( acq_function: AcquisitionFunction, bounds: Tensor, q: int, num_restarts: int, raw_samples: Optional[int] = None, options: Optional[Dict[str, Union[bool, float, int, str]]] = None, inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, nonlinear_inequality_constraints: Optional[List[Callable]] = None, fixed_features: Optional[Dict[int, float]] = None, post_processing_func: Optional[Callable[[Tensor], Tensor]] = None, batch_initial_conditions: Optional[Tensor] = None, return_best_only: bool = True, sequential: bool = False, **kwargs: Any, ) -> Tuple[Tensor, Tensor]: r"""Generate a set of candidates via multi-start optimization. Args: acq_function: An AcquisitionFunction. bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`. q: The number of candidates. num_restarts: The number of starting points for multistart acquisition function optimization. raw_samples: The number of samples for initialization. This is required if `batch_initial_conditions` is not specified. options: Options for candidate generation. inequality constraints: A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form `\sum_i (X[indices[i]] * coefficients[i]) >= rhs` equality constraints: A list of tuples (indices, coefficients, rhs), with each tuple encoding an inequality constraint of the form `\sum_i (X[indices[i]] * coefficients[i]) = rhs` nonlinear_inequality_constraints: A list of callables with that represent non-linear inequality constraints of the form `callable(x) >= 0`. Each callable is expected to take a `(num_restarts) x q x d`-dim tensor as an input and return a `(num_restarts) x q`-dim tensor with the constraint values. The constraints will later be passed to SLSQP. You need to pass in `batch_initial_conditions` in this case. Using non-linear inequality constraints also requires that `batch_limit` is set to 1, which will be done automatically if not specified in `options`. fixed_features: A map `{feature_index: value}` for features that should be fixed to a particular value during generation. post_processing_func: A function that post-processes an optimization result appropriately (i.e., according to `round-trip` transformations). batch_initial_conditions: A tensor to specify the initial conditions. Set this if you do not want to use default initialization strategy. return_best_only: If False, outputs the solutions corresponding to all random restart initializations of the optimization. sequential: If False, uses joint optimization, otherwise uses sequential optimization. kwargs: Additonal keyword arguments. Returns: A two-element tuple containing - a `(num_restarts) x q x d`-dim tensor of generated candidates. - a tensor of associated acquisition values. If `sequential=False`, this is a `(num_restarts)`-dim tensor of joint acquisition values (with explicit restart dimension if `return_best_only=False`). If `sequential=True`, this is a `q`-dim tensor of expected acquisition values conditional on having observed canidates `0,1,...,i-1`. Example: >>> # generate `q=2` candidates jointly using 20 random restarts >>> # and 512 raw samples >>> candidates, acq_value = optimize_acqf(qEI, bounds, 2, 20, 512) >>> generate `q=3` candidates sequentially using 15 random restarts >>> # and 256 raw samples >>> qEI = qExpectedImprovement(model, best_f=0.2) >>> bounds = torch.tensor([[0.], [1.]]) >>> candidates, acq_value_list = optimize_acqf( >>> qEI, bounds, 3, 15, 256, sequential=True >>> ) """ if sequential and q > 1: if not return_best_only: raise NotImplementedError( "`return_best_only=False` only supported for joint optimization." ) if isinstance(acq_function, OneShotAcquisitionFunction): raise NotImplementedError( "sequential optimization currently not supported for one-shot " "acquisition functions. Must have `sequential=False`.") candidate_list, acq_value_list = [], [] base_X_pending = acq_function.X_pending for i in range(q): candidate, acq_value = optimize_acqf( acq_function=acq_function, bounds=bounds, q=1, num_restarts=num_restarts, raw_samples=raw_samples, options=options or {}, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, nonlinear_inequality_constraints= nonlinear_inequality_constraints, fixed_features=fixed_features, post_processing_func=post_processing_func, batch_initial_conditions=None, return_best_only=True, sequential=False, ) candidate_list.append(candidate) acq_value_list.append(acq_value) candidates = torch.cat(candidate_list, dim=-2) acq_function.set_X_pending( torch.cat([base_X_pending, candidates], dim=-2 ) if base_X_pending is not None else candidates) logger.info(f"Generated sequential candidate {i+1} of {q}") # Reset acq_func to previous X_pending state acq_function.set_X_pending(base_X_pending) return candidates, torch.stack(acq_value_list) options = options or {} # Handle the trivial case when all features are fixed if fixed_features is not None and len(fixed_features) == bounds.shape[-1]: X = torch.tensor( [fixed_features[i] for i in range(bounds.shape[-1])], device=bounds.device, dtype=bounds.dtype, ) X = X.expand(q, *X.shape) with torch.no_grad(): acq_value = acq_function(X) return X, acq_value if batch_initial_conditions is None: if nonlinear_inequality_constraints: raise NotImplementedError( "`batch_initial_conditions` must be given if there are non-linear " "inequality constraints.") if raw_samples is None: raise ValueError( "Must specify `raw_samples` when `batch_initial_conditions` is `None`." ) ic_gen = (gen_one_shot_kg_initial_conditions if isinstance( acq_function, qKnowledgeGradient) else gen_batch_initial_conditions) batch_initial_conditions = ic_gen( acq_function=acq_function, bounds=bounds, q=q, num_restarts=num_restarts, raw_samples=raw_samples, fixed_features=fixed_features, options=options, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, ) batch_limit: int = options.get( "batch_limit", num_restarts if not nonlinear_inequality_constraints else 1) batch_candidates_list: List[Tensor] = [] batch_acq_values_list: List[Tensor] = [] batched_ics = batch_initial_conditions.split(batch_limit) for i, batched_ics_ in enumerate(batched_ics): # optimize using random restart optimization batch_candidates_curr, batch_acq_values_curr = gen_candidates_scipy( initial_conditions=batched_ics_, acquisition_function=acq_function, lower_bounds=bounds[0], upper_bounds=bounds[1], options={ k: v for k, v in options.items() if k not in INIT_OPTION_KEYS }, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, nonlinear_inequality_constraints=nonlinear_inequality_constraints, fixed_features=fixed_features, ) batch_candidates_list.append(batch_candidates_curr) batch_acq_values_list.append(batch_acq_values_curr) logger.info(f"Generated candidate batch {i+1} of {len(batched_ics)}.") batch_candidates = torch.cat(batch_candidates_list) batch_acq_values = torch.cat(batch_acq_values_list) if post_processing_func is not None: batch_candidates = post_processing_func(batch_candidates) if return_best_only: best = torch.argmax(batch_acq_values.view(-1), dim=0) batch_candidates = batch_candidates[best] batch_acq_values = batch_acq_values[best] if isinstance(acq_function, OneShotAcquisitionFunction): if not kwargs.get("return_full_tree", False): batch_candidates = acq_function.extract_candidates( X_full=batch_candidates) return batch_candidates, batch_acq_values
def gen( self, model_gen_options: Optional[Dict[str, Any]] = None, explore_features: Optional[List[int]] = None, ) -> Tuple[Tensor, Optional[List[Dict[str, Any]]]]: """Generate candidate by optimizing acquisition function. Args: model_gen_options: Dictionary with options for generating candidate, such as SGD parameters. See code for all options and their defaults. explore_features: List of features that will be selected randomly and then fixed for acquisition fn optimization. Returns: Xopt: (1 x d) tensor of the generated candidate candidate_metadata: List of dict of metadata for each candidate. Contains acquisition value for the candidate. """ # Default optimization settings # TODO are these sufficiently robust? Can they be tuned better? options = model_gen_options or {} num_restarts = options.get("num_restarts", 10) raw_samples = options.get("raw_samples", 1000) verbosity_freq = options.get("verbosity_freq", -1) lr = options.get("lr", 0.01) momentum = options.get("momentum", 0.9) nesterov = options.get("nesterov", True) epochs = options.get("epochs", 50) milestones = options.get("milestones", [25, 40]) gamma = options.get("gamma", 0.1) loss_constraint_fun = options.get( "loss_constraint_fun", default_loss_constraint_fun ) acq_function = self._get_acquisition_fn() # Augment bounds with deriv indicator bounds = torch.cat((self.bounds_, torch.zeros(2, 1, dtype=self.dtype)), dim=1) # Fix deriv indicator to 0 during optimization fixed_features = {(bounds.shape[1] - 1): 0.0} # Fix explore features to random values if explore_features is not None: for idx in explore_features: val = ( bounds[0, idx] + torch.rand(1, dtype=self.dtype) * (bounds[1, idx] - bounds[0, idx]) ).item() fixed_features[idx] = val bounds[0, idx] = val bounds[1, idx] = val # Initialize batch_initial_conditions = gen_batch_initial_conditions( acq_function=acq_function, bounds=bounds, q=1, num_restarts=num_restarts, raw_samples=raw_samples, ) clamped_candidates = columnwise_clamp( X=batch_initial_conditions, lower=bounds[0], upper=bounds[1] ).requires_grad_(True) candidates = fix_features(clamped_candidates, fixed_features) optimizer = torch.optim.SGD( params=[clamped_candidates], lr=lr, momentum=momentum, nesterov=nesterov ) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=milestones, gamma=gamma ) # Optimize for epoch in range(epochs): loss = -acq_function(candidates).sum() # adjust loss based on constraints on candidates loss = loss_constraint_fun(loss, candidates) if verbosity_freq > 0 and epoch % verbosity_freq == 0: logger.info("Iter: {} - Value: {:.3f}".format(epoch, -(loss.item()))) def closure(): optimizer.zero_grad() loss.backward( retain_graph=True ) # Variational model requires retain_graph return loss optimizer.step(closure) clamped_candidates.data = columnwise_clamp( X=clamped_candidates, lower=bounds[0], upper=bounds[1] ) candidates = fix_features(clamped_candidates, fixed_features) lr_scheduler.step() # Extract best point with torch.no_grad(): batch_acquisition = acq_function(candidates) best = torch.argmax(batch_acquisition.view(-1), dim=0) Xopt = candidates[best][:, :-1].detach() candidate_metadata = [{"acquisition_value": batch_acquisition[best].item()}] return Xopt, candidate_metadata