fea: add the option to provide both some initial batch conditions and request some raw samples

Author: CompRhys

botorch/optim/optimize.py

@@ -109,12 +109,31 @@ def __post_init__(self) -> None:
                     "3-dimensional. Its shape is "
                     f"{batch_initial_conditions_shape}."
                 )
+
             if batch_initial_conditions_shape[-1] != d:
                 raise ValueError(
                     f"batch_initial_conditions.shape[-1] must be {d}. The "
                     f"shape is {batch_initial_conditions_shape}."
                 )
 
+            if (
+                self.raw_samples is not None
+                and (self.raw_samples - batch_initial_conditions_shape[-2]) > 0
+                and len(batch_initial_conditions_shape) == 3
+                and self.num_restarts is not None
+                and batch_initial_conditions_shape[0] not in [1, self.num_restarts]
+            ):
+                warnings.warn(
+                    "If using `batch_initial_conditions` together with `raw_samples`, "
+                    "the first repeat dimension of `batch_initial_conditions` must "
+                    "match `num_restarts`. In the future this will raise an error. "
+                    "Defaulting to old behavior of ignoring `raw_samples` by setting "
+                    "it to None.",
+                    DeprecationWarning,
+                )
+                # Use object.__setattr__ to bypass immutability and set a value
+                object.__setattr__(self, "raw_samples", None)
+
         elif self.ic_generator is None:
             if self.nonlinear_inequality_constraints is not None:
                 raise RuntimeError(
@@ -253,27 +272,73 @@ def _optimize_acqf_sequential_q(
     return candidates, torch.stack(acq_value_list)
 
 
+def _combine_initial_conditions(
+    provided_initial_conditions: Tensor | None = None,
+    generated_initial_conditions: Tensor | None = None,
+    num_restarts: int | None = None,
+) -> Tensor:
+
+    if (
+        provided_initial_conditions is not None
+        and generated_initial_conditions is not None
+    ):
+        if (  # Repeat the provided initial conditions to match the number of restarts
+            provided_initial_conditions.shape[0] == 1
+            and num_restarts is not None
+            and num_restarts > 1
+        ):
+            provided_initial_conditions = provided_initial_conditions.repeat(
+                num_restarts, *([1] * (provided_initial_conditions.dim() - 1))
+            )
+        initial_conditions = torch.cat(
+            [provided_initial_conditions, generated_initial_conditions], dim=-2
+        )
+        perm = torch.randperm(
+            initial_conditions.shape[-2], device=initial_conditions.device
+        )
+        return initial_conditions.gather(
+            -2, perm.unsqueeze(-1).expand_as(initial_conditions)
+        )
+    elif provided_initial_conditions is not None:
+        return provided_initial_conditions
+    elif generated_initial_conditions is not None:
+        return generated_initial_conditions
+    else:
+        raise ValueError(
+            "Either `batch_initial_conditions` or `raw_samples` must be set."
+        )
+
+
 def _optimize_acqf_batch(opt_inputs: OptimizeAcqfInputs) -> tuple[Tensor, Tensor]:
     options = opt_inputs.options or {}
 
-    initial_conditions_provided = opt_inputs.batch_initial_conditions is not None
+    required_raw_samples = opt_inputs.raw_samples
+    generated_initial_conditions = None
 
-    if initial_conditions_provided:
-        batch_initial_conditions = opt_inputs.batch_initial_conditions
-    else:
-        # pyre-ignore[28]: Unexpected keyword argument `acq_function` to anonymous call.
-        batch_initial_conditions = opt_inputs.get_ic_generator()(
-            acq_function=opt_inputs.acq_function,
-            bounds=opt_inputs.bounds,
-            q=opt_inputs.q,
-            num_restarts=opt_inputs.num_restarts,
-            raw_samples=opt_inputs.raw_samples,
-            fixed_features=opt_inputs.fixed_features,
-            options=options,
-            inequality_constraints=opt_inputs.inequality_constraints,
-            equality_constraints=opt_inputs.equality_constraints,
-            **opt_inputs.ic_gen_kwargs,
-        )
+    if required_raw_samples is not None:
+        if opt_inputs.batch_initial_conditions is not None:
+            required_raw_samples -= opt_inputs.batch_initial_conditions.shape[-2]
+
+        if required_raw_samples > 0:
+            # pyre-ignore[28]: Unexpected keyword argument `acq_function` to anonymous call.
+            generated_initial_conditions = opt_inputs.get_ic_generator()(
+                acq_function=opt_inputs.acq_function,
+                bounds=opt_inputs.bounds,
+                q=opt_inputs.q,
+                num_restarts=opt_inputs.num_restarts,
+                raw_samples=required_raw_samples,
+                fixed_features=opt_inputs.fixed_features,
+                options=options,
+                inequality_constraints=opt_inputs.inequality_constraints,
+                equality_constraints=opt_inputs.equality_constraints,
+                **opt_inputs.ic_gen_kwargs,
+            )
+
+    batch_initial_conditions = _combine_initial_conditions(
+        provided_initial_conditions=opt_inputs.batch_initial_conditions,
+        generated_initial_conditions=generated_initial_conditions,
+        num_restarts=opt_inputs.num_restarts,
+    )
 
     batch_limit: int = options.get(
         "batch_limit",
@@ -344,31 +409,38 @@ def _optimize_batch_candidates() -> tuple[Tensor, Tensor, list[Warning]]:
         first_warn_msg = (
             "Optimization failed in `gen_candidates_scipy` with the following "
             f"warning(s):\n{[w.message for w in ws]}\nBecause you specified "
-            "`batch_initial_conditions`, optimization will not be retried with "
-            "new initial conditions and will proceed with the current solution."
-            " Suggested remediation: Try again with different "
-            "`batch_initial_conditions`, or don't provide `batch_initial_conditions.`"
-            if initial_conditions_provided
+            "`batch_initial_conditions`>`raw_samples`, optimization will not "
+            "be retried with new initial conditions and will proceed with the "
+            "current solution. Suggested remediation: Try again with different "
+            "`batch_initial_conditions`, don't provide `batch_initial_conditions`, "
+            "or increase `raw_samples`.`"
+            if required_raw_samples is not None and required_raw_samples <= 0
             else "Optimization failed in `gen_candidates_scipy` with the following "
             f"warning(s):\n{[w.message for w in ws]}\nTrying again with a new "
             "set of initial conditions."
         )
         warnings.warn(first_warn_msg, RuntimeWarning, stacklevel=2)
 
-        if not initial_conditions_provided:
-            batch_initial_conditions = opt_inputs.get_ic_generator()(
+        if required_raw_samples is not None and required_raw_samples > 0:
+            generated_initial_conditions = opt_inputs.get_ic_generator()(
                 acq_function=opt_inputs.acq_function,
                 bounds=opt_inputs.bounds,
                 q=opt_inputs.q,
                 num_restarts=opt_inputs.num_restarts,
-                raw_samples=opt_inputs.raw_samples,
+                raw_samples=required_raw_samples,
                 fixed_features=opt_inputs.fixed_features,
                 options=options,
                 inequality_constraints=opt_inputs.inequality_constraints,
                 equality_constraints=opt_inputs.equality_constraints,
                 **opt_inputs.ic_gen_kwargs,
             )
 
+            batch_initial_conditions = _combine_initial_conditions(
+                provided_initial_conditions=opt_inputs.batch_initial_conditions,
+                generated_initial_conditions=generated_initial_conditions,
+                num_restarts=opt_inputs.num_restarts,
+            )
+
             batch_candidates, batch_acq_values, ws = _optimize_batch_candidates()
 
             optimization_warning_raised = any(
@@ -1177,7 +1249,7 @@ def _gen_batch_initial_conditions_local_search(
     inequality_constraints: list[tuple[Tensor, Tensor, float]],
     min_points: int,
     max_tries: int = 100,
-):
+) -> Tensor:
     """Generate initial conditions for local search."""
     device = discrete_choices[0].device
     dtype = discrete_choices[0].dtype
@@ -1197,6 +1269,66 @@ def _gen_batch_initial_conditions_local_search(
     raise RuntimeError(f"Failed to generate at least {min_points} initial conditions")
 
 
+def _gen_starting_points_local_search(
+    discrete_choices: list[Tensor],
+    raw_samples: int,
+    batch_initial_conditions: Tensor,
+    X_avoid: Tensor,
+    inequality_constraints: list[tuple[Tensor, Tensor, float]],
+    min_points: int,
+    acq_function: AcquisitionFunction,
+    max_batch_size: int = 2048,
+    max_tries: int = 100,
+) -> Tensor:
+    required_min_points = min_points
+    provided_X0 = None
+    generated_X0 = None
+
+    if batch_initial_conditions is not None:
+        provided_X0 = _filter_invalid(
+            X=batch_initial_conditions.squeeze(1), X_avoid=X_avoid
+        )
+        provided_X0 = _filter_infeasible(
+            X=provided_X0, inequality_constraints=inequality_constraints
+        ).unsqueeze(1)
+        required_min_points -= batch_initial_conditions.shape[0]
+
+    if required_min_points > 0:
+        generated_X0 = _gen_batch_initial_conditions_local_search(
+            discrete_choices=discrete_choices,
+            raw_samples=raw_samples,
+            X_avoid=X_avoid,
+            inequality_constraints=inequality_constraints,
+            min_points=min_points,
+            max_tries=max_tries,
+        )
+
+        # pick the best starting points
+        with torch.no_grad():
+            acqvals_init = _split_batch_eval_acqf(
+                acq_function=acq_function,
+                X=generated_X0.unsqueeze(1),
+                max_batch_size=max_batch_size,
+            ).unsqueeze(-1)
+
+        generated_X0 = generated_X0[
+            acqvals_init.topk(k=min_points, largest=True, dim=0).indices
+        ]
+
+    if provided_X0 is not None and generated_X0 is not None:
+        X0 = torch.cat([provided_X0, generated_X0], dim=0)
+    elif provided_X0 is not None:
+        X0 = provided_X0
+    elif generated_X0 is not None:
+        X0 = generated_X0
+    else:
+        raise ValueError(
+            "Either `batch_initial_conditions` or `raw_samples` must be set."
+        )
+
+    return X0
+
+
 def optimize_acqf_discrete_local_search(
     acq_function: AcquisitionFunction,
     discrete_choices: list[Tensor],
@@ -1207,6 +1339,7 @@ def optimize_acqf_discrete_local_search(
     X_avoid: Tensor | None = None,
     batch_initial_conditions: Tensor | None = None,
     max_batch_size: int = 2048,
+    max_tries: int = 100,
     unique: bool = True,
 ) -> tuple[Tensor, Tensor]:
     r"""Optimize acquisition function over a lattice.
@@ -1238,6 +1371,8 @@ def optimize_acqf_discrete_local_search(
         max_batch_size: The maximum number of choices to evaluate in batch.
             A large limit can cause excessive memory usage if the model has
             a large training set.
+        max_tries: Maximum number of iterations to try when generating initial
+            conditions.
         unique: If True return unique choices, o/w choices may be repeated
             (only relevant if `q > 1`).
 
@@ -1247,6 +1382,13 @@ def optimize_acqf_discrete_local_search(
         - a `q x d`-dim tensor of generated candidates.
         - an associated acquisition value.
     """
+    if batch_initial_conditions is not None:
+        if not (
+            len(batch_initial_conditions.shape) == 3
+            and batch_initial_conditions.shape[-2] == 1
+        ):
+            raise ValueError("batch_initial_conditions must have shape `n x 1 x d` if given.")
+
     candidate_list = []
     base_X_pending = acq_function.X_pending if q > 1 else None
     base_X_avoid = X_avoid
@@ -1259,27 +1401,18 @@ def optimize_acqf_discrete_local_search(
     inequality_constraints = inequality_constraints or []
     for i in range(q):
         # generate some starting points
-        if i == 0 and batch_initial_conditions is not None:
-            X0 = _filter_invalid(X=batch_initial_conditions.squeeze(1), X_avoid=X_avoid)
-            X0 = _filter_infeasible(
-                X=X0, inequality_constraints=inequality_constraints
-            ).unsqueeze(1)
-        else:
-            X_init = _gen_batch_initial_conditions_local_search(
-                discrete_choices=discrete_choices,
-                raw_samples=raw_samples,
-                X_avoid=X_avoid,
-                inequality_constraints=inequality_constraints,
-                min_points=num_restarts,
-            )
-            # pick the best starting points
-            with torch.no_grad():
-                acqvals_init = _split_batch_eval_acqf(
-                    acq_function=acq_function,
-                    X=X_init.unsqueeze(1),
-                    max_batch_size=max_batch_size,
-                ).unsqueeze(-1)
-            X0 = X_init[acqvals_init.topk(k=num_restarts, largest=True, dim=0).indices]
+        X0 = _gen_starting_points_local_search(
+            discrete_choices=discrete_choices,
+            raw_samples=raw_samples,
+            batch_initial_conditions=batch_initial_conditions,
+            X_avoid=X_avoid,
+            inequality_constraints=inequality_constraints,
+            min_points=num_restarts,
+            acq_function=acq_function,
+            max_batch_size=max_batch_size,
+            max_tries=max_tries,
+        )
+        batch_initial_conditions = None
 
         # optimize from the best starting points
         best_xs = torch.zeros(len(X0), dim, device=device, dtype=dtype)

botorch/optim/optimize_homotopy.py

@@ -157,7 +157,6 @@ def optimize_acqf_homotopy(
     """
     shared_optimize_acqf_kwargs = {
         "num_restarts": num_restarts,
-        "raw_samples": raw_samples,
         "inequality_constraints": inequality_constraints,
         "equality_constraints": equality_constraints,
         "nonlinear_inequality_constraints": nonlinear_inequality_constraints,
@@ -181,11 +180,14 @@ def optimize_acqf_homotopy(
         homotopy.restart()
 
         while not homotopy.should_stop:
+            # After the first iteration we don't want to generate new raw samples
+            requested_raw_samples = raw_samples if candidates is None else None
             candidates, acq_values = optimize_acqf(
                 acq_function=acq_function,
                 bounds=bounds,
                 q=1,
                 options=options,
+                raw_samples=requested_raw_samples,
                 batch_initial_conditions=candidates,
                 **shared_optimize_acqf_kwargs,
             )
@@ -204,6 +206,7 @@ def optimize_acqf_homotopy(
             bounds=bounds,
             q=1,
             options=final_options,
+            raw_samples=None,
             batch_initial_conditions=candidates,
             **shared_optimize_acqf_kwargs,
         )

botorch/optim/optimize_mixed.py

@@ -496,6 +496,7 @@ def continuous_step(
     updated_opt_inputs = dataclasses.replace(
         opt_inputs,
         q=1,
+        raw_samples=None,
         num_restarts=1,
         batch_initial_conditions=current_x.unsqueeze(0),
         fixed_features={

botorch/posteriors/posterior.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod, abstractproperty
+from abc import ABC, abstractmethod
 
 import torch
 from torch import Tensor
@@ -77,12 +77,14 @@ def sample(self, sample_shape: torch.Size | None = None) -> Tensor:
         with torch.no_grad():
             return self.rsample(sample_shape=sample_shape)
 
-    @abstractproperty
+    @property
+    @abstractmethod
     def device(self) -> torch.device:
         r"""The torch device of the distribution."""
         pass  # pragma: no cover
 
-    @abstractproperty
+    @property
+    @abstractmethod
     def dtype(self) -> torch.dtype:
         r"""The torch dtype of the distribution."""
         pass  # pragma: no cover