diff --git a/bax_algorithms/pathwise/optimize.py b/bax_algorithms/pathwise/optimize.py
index e7b5686..2f2639a 100644
--- a/bax_algorithms/pathwise/optimize.py
+++ b/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))
diff --git a/bax_algorithms/utils.py b/bax_algorithms/utils.py
index 0902992..1d16d2c 100644
--- a/bax_algorithms/utils.py
+++ b/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
diff --git a/examples/solenoid_alignment_constraints.ipynb b/examples/solenoid_alignment_constraints.ipynb
new file mode 100644
index 0000000..9be683e
--- /dev/null
+++ b/examples/solenoid_alignment_constraints.ipynb
@@ -0,0 +1,507 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "26acb126",
+ "metadata": {},
+ "source": [
+ "# BAX final prediction with an artificial 2D charge constraint\n",
+ "\n",
+ "This MWE compares `get_bax_mean_prediction()` with the general\n",
+ "`get_constrained_mean_prediction()` helper in `bax_algorithms.utils`.\n",
+ "Both optimize the same virtual alignment objective using fitted GP means.\n",
+ "\n",
+ "The artificial charge depends **only on the two tuning coordinates, x and y**.\n",
+ "A small elliptical region near good alignment has insufficient charge.\n",
+ "The field is used only for the virtual alignment measurement; it does not enter\n",
+ "the charge function or the definition of the excluded region.\n",
+ "\n",
+ "Run all cells with this repository and its dependencies installed. The notebook\n",
+ "works when started in either the repository root or `examples/`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a0ff5c33",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-09-22T20:11:03.003103Z",
+ "iopub.status.busy": "2026-09-22T20:11:03.002961Z",
+ "iopub.status.idle": "2026-09-22T20:11:04.075066Z",
+ "shell.execute_reply": "2026-09-22T20:11:04.074744Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import matplotlib.pyplot as plt\n",
+ "from matplotlib.lines import Line2D\n",
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import torch\n",
+ "from IPython.display import display\n",
+ "from xopt import Xopt\n",
+ "from xopt.evaluator import Evaluator\n",
+ "from xopt.generators.bayesian.bax_generator import BaxGenerator\n",
+ "from xopt.vocs import VOCS\n",
+ "\n",
+ "from bax_algorithms.pathwise.optimize import DifferentialEvolution\n",
+ "from bax_algorithms.solenoid_alignment import PathwiseSolenoidAlignment\n",
+ "from bax_algorithms.utils import (\n",
+ " get_bax_mean_prediction,\n",
+ " get_constrained_mean_prediction,\n",
+ ")\n",
+ "\n",
+ "torch.manual_seed(1)\n",
+ "np.random.seed(1)\n",
+ "torch.set_num_threads(1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cf62416e",
+ "metadata": {},
+ "source": [
+ "## Synthetic measurement and charge\n",
+ "\n",
+ "The centroids are $c_x=(x-0.1)b$ and $c_y=1.3(y+0.05)b$, where $b$ is the\n",
+ "measurement field. The alignment objective is the sum of their standard\n",
+ "deviations over a field scan; its true minimum is at $(0.1,-0.05)$.\n",
+ "\n",
+ "Artificial charge in pC is\n",
+ "\n",
+ "$$Q(x,y)=100-80\\exp\\left[-\\frac12\\left(\n",
+ "\\frac{(x-0.18)^2}{0.22^2}+\\frac{(y+0.05)^2}{0.25^2}\\right)\\right].$$\n",
+ "\n",
+ "We require $Q\\ge60$ pC. The excluded ellipse occupies about 6% of the tuning\n",
+ "square and includes the unconstrained optimum. Nothing in $Q$ depends on field."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "a2e3437a",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-09-22T20:11:04.076316Z",
+ "iopub.status.busy": "2026-09-22T20:11:04.076234Z",
+ "iopub.status.idle": "2026-09-22T20:11:04.083115Z",
+ "shell.execute_reply": "2026-09-22T20:11:04.082754Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Training observations: 363\n",
+ "Model input order: ['x', 'y', 'field']\n"
+ ]
+ }
+ ],
+ "source": [
+ "CHARGE_THRESHOLD = 60.0\n",
+ "\n",
+ "\n",
+ "def charge(x, y):\n",
+ " radius_squared = ((x - 0.18) / 0.22) ** 2 + ((y + 0.05) / 0.25) ** 2\n",
+ " return 100.0 - 80.0 * np.exp(-0.5 * radius_squared)\n",
+ "\n",
+ "\n",
+ "def measure(inputs):\n",
+ " x, y, field = inputs[\"x\"], inputs[\"y\"], inputs[\"field\"]\n",
+ " return {\n",
+ " \"cx\": (x - 0.1) * field,\n",
+ " \"cy\": 1.3 * (y + 0.05) * field,\n",
+ " \"charge\": charge(x, y),\n",
+ " }\n",
+ "\n",
+ "\n",
+ "vocs = VOCS(\n",
+ " variables={\"x\": [-1.0, 1.0], \"y\": [-1.0, 1.0], \"field\": [-1.0, 1.0]},\n",
+ " observables=[\"cx\", \"cy\"],\n",
+ " constraints={\"charge\": [\"GREATER_THAN\", CHARGE_THRESHOLD]},\n",
+ ")\n",
+ "algorithm = PathwiseSolenoidAlignment(\n",
+ " x_key=\"cx\",\n",
+ " y_key=\"cy\",\n",
+ " meas_dim=vocs.variable_names.index(\"field\"),\n",
+ " observable_names_ordered=[\"cx\", \"cy\"],\n",
+ " n_steps_measurement_param=5,\n",
+ " n_samples=2,\n",
+ ")\n",
+ "generator = BaxGenerator(vocs=vocs, algorithm=algorithm)\n",
+ "generator.gp_constructor.use_low_noise_prior = True\n",
+ "X = Xopt(vocs=vocs, generator=generator, evaluator=Evaluator(function=measure))\n",
+ "\n",
+ "# A deterministic grid supplies training data; no acquisition loop is needed\n",
+ "# to isolate the behavior of the final-prediction helpers.\n",
+ "training_inputs = pd.DataFrame(\n",
+ " [\n",
+ " {\"x\": x, \"y\": y, \"field\": field}\n",
+ " for x in np.linspace(-1, 1, 11)\n",
+ " for y in np.linspace(-1, 1, 11)\n",
+ " for field in [-1.0, 0.0, 1.0]\n",
+ " ]\n",
+ ")\n",
+ "X.evaluate_data(training_inputs)\n",
+ "print(f\"Training observations: {len(X.data)}\")\n",
+ "print(\"Model input order:\", vocs.variable_names)\n",
+ "assert (\n",
+ " measure({\"x\": 0.1, \"y\": -0.05, \"field\": -1})[\"charge\"]\n",
+ " == measure({\"x\": 0.1, \"y\": -0.05, \"field\": 1})[\"charge\"]\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "68b41667",
+ "metadata": {},
+ "source": [
+ "## Compare the two final predictions\n",
+ "\n",
+ "`get_constrained_mean_prediction()` reads every constraint from VOCS, with Xopt's\n",
+ "output ordering and inequality directions. No charge name, alignment algorithm,\n",
+ "or field scan is hard-coded into that helper.\n",
+ "\n",
+ "Here `kappa=2` requires the charge GP's lower confidence bound\n",
+ "`mean - 2 * std >= 60`. Use `kappa=0` for posterior-mean feasibility alone.\n",
+ "This margin is not a guarantee of physical safety.\n",
+ "\n",
+ "As with the existing optimizer, dimensions outside the tuning subset are filled\n",
+ "with their bounds' midpoints for the constraint query. Here that is `field=0`;\n",
+ "the true constraint is independent of field. A general `constraint_points`\n",
+ "callback is available if an application needs checks at other inputs."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "b92bbf0a",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-09-22T20:11:04.084071Z",
+ "iopub.status.busy": "2026-09-22T20:11:04.084022Z",
+ "iopub.status.idle": "2026-09-22T20:11:05.320083Z",
+ "shell.execute_reply": "2026-09-22T20:11:05.319711Z"
+ }
+ },
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " x | \n",
+ " y | \n",
+ " predicted misalignment | \n",
+ " charge mean (pC) | \n",
+ " charge lower bound (pC) | \n",
+ " true charge (pC) | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | Existing helper | \n",
+ " 0.10040 | \n",
+ " -0.05021 | \n",
+ " 0.00000 | \n",
+ " 25.22454 | \n",
+ " 25.05568 | \n",
+ " 25.06837 | \n",
+ "
\n",
+ " \n",
+ " | Constraint-aware helper | \n",
+ " -0.08017 | \n",
+ " -0.05026 | \n",
+ " 0.15414 | \n",
+ " 60.17362 | \n",
+ " 60.00706 | \n",
+ " 60.24339 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "