Add function for constraint optimization for solenoid alignment, and a MWE notebook

bax_algorithms/pathwise/optimize.py

@@ -7,7 +7,8 @@
 import torch
 import numpy as np
 import time
-from scipy.optimize import differential_evolution
+import warnings
+from scipy.optimize import NonlinearConstraint, differential_evolution
 
 
 class VirtualOptimizer(XoptBaseModel, ABC):
@@ -27,9 +28,14 @@ def optimize(
         bounds: Tensor,
         n_samples: int,
         optimization_indeces: Tensor = None,
+        constraints: Callable | None = None,
     ) -> Tensor:
         """
-        Minimizes virtual objective sample functions and returns optimal inputs.
+        Optimizes virtual objective sample functions and returns optimal inputs.
+
+        constraints, when provided, maps full inputs of shape
+        (n_samples, n_points, n_dims) to residuals of shape
+        (n_samples, n_points, n_constraints). Every residual must be <= 0.
         """
         raise NotImplementedError("This method should be implemented in subclasses.")
 
@@ -104,6 +110,7 @@ def optimize(
         bounds: Tensor,
         n_samples: int,
         optimization_indeces: Tensor = None,
+        constraints: Callable | None = None,
     ) -> Tensor:
         target_function = self._get_target_function(
             virtual_objective,
@@ -117,13 +124,35 @@ def optimize(
             bounds, n_samples, optimization_indeces
         )
 
+        def constraint_function(x):
+            full_inputs = self._prepare_inputs(
+                x, bounds, n_samples, optimization_indeces
+            )
+            residuals = constraints(full_inputs)
+            if residuals.ndim != 3 or residuals.shape[:2] != full_inputs.shape[:2]:
+                raise ValueError(
+                    "constraints must return (n_samples, n_points, n_constraints)"
+                )
+            # Keep every sample's constraints separate rather than summing them.
+            residuals = residuals.permute(0, 2, 1).reshape(-1, full_inputs.shape[1])
+            residuals = residuals.detach().cpu().numpy()
+            residuals = np.where(np.isfinite(residuals), residuals, np.inf)
+            return residuals[:, 0] if x.ndim == 1 else residuals
+
+        scipy_constraints = (
+            (NonlinearConstraint(constraint_function, -np.inf, 0.0),)
+            if constraints is not None
+            else ()
+        )
         if self.verbose:
             start = time.time()
             print("Beginning BAX virtual optimization.")
         res = differential_evolution(
             target_function,
             bounds=de_bounds,
+            constraints=scipy_constraints,
             vectorized=True,
+            updating="deferred",
             polish=self.polish,
             popsize=self.popsize,
             maxiter=self.maxiter,
@@ -132,15 +161,35 @@ def optimize(
         if self.verbose:
             print("BAX virtual optimization took:", time.time() - start, "seconds.")
 
+        if constraints is not None:
+            if not np.isfinite(res.fun) or not np.all(constraint_function(res.x) <= 0):
+                raise RuntimeError("No feasible prediction found; do not use result.x")
+            if not res.success:
+                warnings.warn(
+                    f"Feasible result, but optimizer did not converge: {res.message}",
+                    stacklevel=2,
+                )
+
         if optimization_indeces is not None:
             ndim = len(optimization_indeces)
         else:
             ndim = bounds.shape[-1]
 
-        best_x = torch.from_numpy(res.x).reshape(n_samples, 1, ndim)
+        best_x = torch.from_numpy(res.x).to(bounds).reshape(n_samples, 1, ndim)
 
         return best_x
 
+    @staticmethod
+    def _prepare_inputs(x, bounds, n_samples, optimization_indeces):
+        x = torch.from_numpy(x).to(bounds)
+        dim = int(x.shape[0] / n_samples)
+        x = x.reshape(n_samples, dim, -1).transpose(1, 2)
+        if optimization_indeces is None:
+            return x
+        full_inputs = bounds.mean(dim=0).repeat(*x.shape[:-1], 1)
+        full_inputs[..., optimization_indeces] = x
+        return full_inputs
+
     def _wrap_virtual_objective(
         self,
         virtual_objective: Callable,
@@ -150,19 +199,13 @@ def _wrap_virtual_objective(
         optimization_indeces: Tensor,
     ) -> Callable:
         def wrapped_virtual_objective(x):
-            x = torch.from_numpy(x)
-            # x.shape = num_params x popsize*num_params // num_params = dim*n_samples
-            dim = int(x.shape[0] / n_samples)
-            x = x.reshape(n_samples, dim, -1)
-            x = torch.swapaxes(x, 1, 2)
-            if optimization_indeces is not None:
-                x_p = bounds.mean(dim=0).repeat(*x.shape[:-1], 1)
-                x_p[..., optimization_indeces] = x
-            else:
-                x_p = x
+            # SciPy can pass no points when the whole population is infeasible.
+            if x.ndim == 2 and x.shape[1] == 0:
+                return np.empty(0)
+            x_p = self._prepare_inputs(x, bounds, n_samples, optimization_indeces)
             res = virtual_objective(sample_functions_list, x_p, bounds)
-            res = res.sum(dim=0).detach().numpy()
-            return res
+            res = res.sum(dim=0).detach().cpu().numpy()
+            return res.item() if x.ndim == 1 else res
 
         return wrapped_virtual_objective
 
@@ -173,6 +216,9 @@ def _get_virtual_optimization_bounds(
         Get bounds for virtual optimization (may not be the same as bounds passed to optimizer).
         """
         if optimization_indeces is not None:
-            return np.tile(bounds[..., optimization_indeces].numpy().T, (n_samples, 1))
+            return np.tile(
+                bounds[..., optimization_indeces].detach().cpu().numpy().T,
+                (n_samples, 1),
+            )
         else:
-            return np.tile(bounds.numpy().T, (n_samples, 1))
+            return np.tile(bounds.detach().cpu().numpy().T, (n_samples, 1))

bax_algorithms/utils.py

@@ -1,4 +1,8 @@
+from collections.abc import Callable
+import math
+
 from xopt.generators.bayesian.bax_generator import BaxGenerator
+from xopt.generators.bayesian.objectives import create_constraint_callables
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.model_list_gp_regression import ModelListGP
 from bax_algorithms.pathwise.optimize import VirtualOptimizer
@@ -45,6 +49,94 @@ def func(x):
     return best_inputs.squeeze(0)
 
 
+def get_constrained_mean_prediction(
+    generator: BaxGenerator,
+    mean_optimizer: VirtualOptimizer,
+    *,
+    kappa: float = 0.0,
+    constraint_points: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+    | None = None,
+):
+    """Optimize the BAX mean objective subject to all VOCS output constraints.
+
+    Uses Xopt's constraint callables, including output ordering and LESS_THAN /
+    GREATER_THAN directions. kappa=0 requires posterior-mean feasibility. Positive
+    kappa requires the signed constraint mean plus kappa times its posterior
+    standard deviation to be <= 0. This is a pointwise uncertainty margin, not
+    a guarantee of physical safety or joint feasibility.
+
+    By default, constraints are evaluated at the full candidate inputs passed to
+    the virtual objective. Dimensions not optimized by the algorithm are filled
+    with the midpoint of their bounds, as in get_bax_mean_prediction. To enforce
+    constraints at other inputs, provide constraint_points(x, bounds), returning
+    shape (*x.shape[:-1], n_evaluation_points, x.shape[-1]). All evaluation points
+    must satisfy every constraint. This callback can also select a reference
+    setting for dimensions not optimized by the algorithm.
+
+    The optimizer must implement the optional constraints callback in
+    VirtualOptimizer.optimize (implemented by DifferentialEvolution). Without
+    VOCS constraints this delegates to get_bax_mean_prediction unchanged.
+    Returns shape (1, n_tuning_dims), in the algorithm's optimization order.
+    Like the original helper, this expects bounds covering all model inputs;
+    reduced bounds from fixed_features are not supported.
+    """
+    if not math.isfinite(kappa) or kappa < 0:
+        raise ValueError("kappa must be finite and nonnegative")
+    if not generator.vocs.constraints:
+        return get_bax_mean_prediction(generator, mean_optimizer)
+
+    # Retain the constraint models separately from the objective observables.
+    model = generator.train_model()
+    bounds = generator._get_optimization_bounds()
+    if bounds.shape[-1] != len(generator.model_input_names):
+        raise ValueError("Mean prediction requires bounds covering all model inputs")
+    output_names = generator.vocs.output_names
+    observable_models = [
+        model.subset_output([output_names.index(name)])
+        for name in generator.algorithm.observable_names_ordered
+    ]
+    mean_functions = [(lambda x, m=m: m.posterior(x).mean) for m in observable_models]
+    constraint_functions = create_constraint_callables(generator.vocs)
+    constraint_indices = [
+        output_names.index(name) for name in generator.vocs.constraint_names
+    ]
+
+    @torch.no_grad()
+    def constraints(x):
+        points = (
+            x.unsqueeze(-2)
+            if constraint_points is None
+            else constraint_points(x, bounds)
+        )
+        if (
+            points.ndim != x.ndim + 1
+            or points.shape[:-2] != x.shape[:-1]
+            or points.shape[-1] != x.shape[-1]
+            or points.shape[-2] == 0
+        ):
+            raise ValueError(
+                "constraint_points must return (*x.shape[:-1], n_points, n_dims)"
+            )
+        posterior = model.posterior(points)
+        mean = posterior.mean
+        residuals = torch.stack([fn(mean) for fn in constraint_functions], dim=-1)
+        if kappa:
+            std = posterior.variance[..., constraint_indices].clamp_min(0).sqrt()
+            residuals = residuals + kappa * std
+        residuals = torch.where(torch.isfinite(residuals), residuals, torch.inf)
+        return residuals.amax(dim=-2)
+
+    best_inputs = mean_optimizer.optimize(
+        virtual_objective=generator.algorithm.evaluate_virtual_objective,
+        sample_functions_list=mean_functions,
+        bounds=bounds,
+        optimization_indeces=generator.algorithm._get_optimization_indeces(bounds),
+        n_samples=1,
+        constraints=constraints,
+    )
+    return best_inputs.squeeze(0)
+
+
 def tuning_input_tensor_to_dict(generator, x_tuning):
     """
     Converts a single set of tuning parameters to a dictionary for input to Xopt