"""Agent interface for optimization with Ax as the backend."""
import logging
from collections.abc import Mapping, Sequence
from typing import Any, Generic, cast
import bluesky.preprocessors as bpp
from ax import Client, TOutcome, TParameterization
from ax.analysis.plotly.surface.contour import ContourPlot
from ax.core.analysis_card import AnalysisCardBase
from ax.core.types import TParamValue
from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
from bluesky.callbacks import CallbackBase
from bluesky.utils import MsgGenerator
from ..callbacks.logger import OptimizationLogger
from ..callbacks.router import OptimizationCallbackRouter
from ..plan_stubs import navigate_to_best
from ..plans import acquire_baseline, optimize, sample_suggestions
from ..protocols import (
AcquisitionPlan,
Actuator,
EvaluationFunction,
OptimizationProblem,
Sensor,
TUid,
)
from ..utils import InferredReadable
from .dof import DOF, DOFConstraint
from .objective import Objective, OutcomeConstraint, ScalarizedObjective, to_ax_objective_str
from .optimizer import AxOptimizer
logger = logging.getLogger(__name__)
class _AxAgentMixin:
"""
Mixin providing Ax-related functionality shared by agents.
Expects subclasses to define `self._optimizer` as an `AxOptimizer`.
"""
_optimizer: AxOptimizer
@property
def ax_client(self) -> Client:
return self._optimizer.ax_client
@property
def checkpoint_path(self) -> str | None:
return self._optimizer.checkpoint_path
@checkpoint_path.setter
def checkpoint_path(self, value: str) -> None:
self._optimizer.checkpoint_path = value
@property
def fixed_dofs(self) -> dict[str, Any] | None:
return self._optimizer.fixed_parameters
@fixed_dofs.setter
def fixed_dofs(self, fixed_dofs: dict[DOF, Any] | None) -> None:
"""
Fix degrees of freedom to a certain value for future optimizations.
Parameters
----------
fixed_dofs : dict[DOF, Any] | None
A mapping of DOFs to the values they should be fixed to.
"""
if not fixed_dofs:
self._optimizer.fixed_parameters = None
return
self._optimizer.fixed_parameters = {dof.parameter_name: value for dof, value in fixed_dofs.items()}
def suggest(self, num_points: int = 1) -> Sequence[Mapping]:
"""
Get the next point(s) to evaluate in the search space.
Uses the Bayesian optimization algorithm to suggest promising points based
on all previously acquired data. Each suggestion includes an "_id" key for
tracking.
Parameters
----------
num_points : int, optional
The number of points to suggest. Default is 1. Higher values enable
batch optimization but may reduce optimization efficiency per iteration.
Returns
-------
Sequence[Mapping]
A sequence of mappings, each containing a parameterization of a point to
evaluate next. Each mapping includes an "_id" key for identification.
"""
return self._optimizer.suggest(num_points)
def ingest(self, points: Sequence[Mapping]) -> None:
"""
Ingest evaluation results into the optimizer.
Updates the optimizer's model with new data. Can ingest both suggested points
(with "_id" key) and external data (without "_id" key).
Parameters
----------
points : Sequence[Mapping]
A sequence of mappings containing outcomes for a trial. For suggested points,
include the "_id" key. For external data, include DOF names and objective
values, and omit "_id".
Notes
-----
This method is typically called automatically by :meth:`optimize`. Manual usage
is only needed for custom workflows or when ingesting external data.
For complete examples, see :doc:`/how-to-guides/attach-data-to-experiments`.
"""
self._optimizer.ingest(points)
def get_best_points(self) -> Sequence[tuple[int, TParameterization, TOutcome]]:
"""
Get a list of the optimal points found during optimization.
For single-objective optimization, returns a single best point.
For multi-objective optimization, returns the Pareto-optimal set.
Returns
-------
Sequence[tuple[int, TParameterization, TOutcome]]
Each element in the list is a tuple of:
- trial index (int)
- parameter values (dict)
- metric values (dict, where values may be (value, sem) tuples)
See Also
--------
navigate_to_best : Plan stub to move actuators to a best point.
"""
return self._optimizer.get_best_points()
def plot_objective(
self, x_dof_name: str, y_dof_name: str, objective_name: str, *args: Any, **kwargs: Any
) -> list[AnalysisCardBase]:
"""
Plot the predicted objective as a function of two DOFs.
Creates a contour plot showing the model's prediction of an objective across
the space defined by two DOFs. Useful for visualizing the optimization landscape.
Parameters
----------
x_dof_name : str
The name of the DOF to plot on the x-axis.
y_dof_name : str
The name of the DOF to plot on the y-axis.
objective_name : str
The name of the objective to plot.
*args : Any
Additional positional arguments passed to Ax's compute_analyses.
**kwargs : Any
Additional keyword arguments passed to Ax's compute_analyses.
Returns
-------
list[AnalysisCard]
The computed analysis cards containing the plot data.
See Also
--------
ax.analysis.ContourPlot : Pre-built analysis for plotting objectives.
ax.analysis.AnalysisCard : Contains the raw and computed data.
"""
return self.ax_client.compute_analyses(
[
ContourPlot(
x_parameter_name=x_dof_name,
y_parameter_name=y_dof_name,
metric_name=objective_name,
),
],
*args,
**kwargs,
)
def checkpoint(self) -> None:
"""Save the agent's state to a JSON file."""
self._optimizer.checkpoint()
def reconfigure_search_space(self, dof_mappings: dict[DOF, tuple[float, float] | list[TParamValue]]) -> None:
"""
Update bounds or values of existing DOFs for future optimizations.
Parameters
----------
dof_mappings : dict[DOF, tuple[float, float] | list[float] | list[int] | list[str] | list[bool]]
Mapping of DOFs to their new search space.
"""
self._optimizer.reconfigure_search_space({dof.parameter_name: update for dof, update in dof_mappings.items()})
[docs]
class Agent(_AxAgentMixin, Generic[TUid]):
"""
An interface that uses Ax as the backend for optimization and experiment tracking.
The Agent is the main entry point for setting up and running Bayesian optimization
using Bluesky Adaptive. It coordinates the DOFs, objectives, evaluation function, and optimizer
to perform intelligent exploration of the parameter space.
Parameters
----------
sensors : Sequence[Sensor]
The sensors to use for acquisition. These should be the minimal set
of sensors that are needed to compute the objectives.
dofs : Sequence[DOF]
The degrees of freedom that the agent can control, which determine the search space.
objectives : Sequence[Objective]
The objectives which the agent will try to optimize.
evaluation_function : EvaluationFunction
The function to evaluate acquired data and produce outcomes.
acquisition_plan : AcquisitionPlan | None, optional
The acquisition plan to use for acquiring data from the beamline. If not provided,
:func:`bluesky_adaptive.plans.default_acquire` will be used.
dof_constraints : Sequence[DOFConstraint] | None, optional
Constraints on DOFs to refine the search space.
outcome_constraints : Sequence[OutcomeConstraint] | None, optional
Constraints on outcomes to be satisfied during optimization.
checkpoint_path : str | None, optional
The path to the checkpoint file to save the optimizer's state to.
stopping_strategy : BaseGlobalStoppingStrategy | None, optional
A stopping strategy that decides when/if optimization should halt early.
**kwargs : Any
Additional keyword arguments to configure the Ax experiment.
See Also
--------
bluesky_adaptive.protocols.Sensor : The protocol for sensors.
bluesky_adaptive.ax.dof.RangeDOF : For continuous parameters.
bluesky_adaptive.ax.dof.ChoiceDOF : For discrete parameters.
bluesky_adaptive.ax.objective.Objective : For defining objectives.
bluesky_adaptive.ax.optimizer.AxOptimizer : The optimizer used internally.
bluesky_adaptive.plans.optimize : Bluesky plan for running optimization.
Notes
-----
For more complex setups, you can configure the Ax client directly via ``self.ax_client``.
For complete working examples of creating and using an Agent, see the tutorial
documentation, particularly :doc:`/tutorials/simple-experiment`.
"""
def __init__(
self,
sensors: Sequence[Sensor],
dofs: Sequence[DOF],
objectives: Sequence[Objective] | ScalarizedObjective,
evaluation_function: EvaluationFunction[TUid],
acquisition_plan: AcquisitionPlan[TUid] | None = None,
dof_constraints: Sequence[DOFConstraint] | None = None,
outcome_constraints: Sequence[OutcomeConstraint] | None = None,
checkpoint_path: str | None = None,
stopping_strategy: BaseGlobalStoppingStrategy | None = None,
**kwargs: Any,
):
if any(isinstance(dof.actuator, str) for dof in dofs):
dof_actuator_strs = [dof.actuator for dof in dofs if isinstance(dof.actuator, str)]
raise ValueError(
f"DOFs with actuators must be `Actuator` instances, not strings. Got strings for: {dof_actuator_strs}"
)
self._sensors = sensors
self._actuators: Sequence[Actuator] = [cast(Actuator, dof.actuator) for dof in dofs if dof.actuator is not None]
self._evaluation_function = evaluation_function
self._acquisition_plan = acquisition_plan
self._optimizer = AxOptimizer(
parameters=[dof.to_ax_parameter_config() for dof in dofs],
objective=to_ax_objective_str(objectives),
parameter_constraints=[constraint.ax_constraint for constraint in dof_constraints] if dof_constraints else None,
outcome_constraints=[constraint.ax_constraint for constraint in outcome_constraints]
if outcome_constraints
else None,
checkpoint_path=checkpoint_path,
stopping_strategy=stopping_strategy,
**kwargs,
)
self._readable_cache: dict[str, InferredReadable] = {}
self._callbacks: list[CallbackBase] = [OptimizationLogger()]
self._callback_router = OptimizationCallbackRouter(self._callbacks)
[docs]
@classmethod
def from_checkpoint(
cls,
checkpoint_path: str,
actuators: Sequence[Actuator],
sensors: Sequence[Sensor],
evaluation_function: EvaluationFunction[TUid],
acquisition_plan: AcquisitionPlan[TUid] | None = None,
) -> "Agent[TUid]":
"""
Load an agent from the optimizer's checkpoint file.
.. note::
Only the optimizer state is saved during a checkpoint, so we cannot reliably validate
the remaining state against the optimizer configuration.
Parameters
----------
checkpoint_path : str
The checkpoint path to load the agent from.
actuators: Sequence[Actuator]
Objects that can be moved to control the beamline using the Bluesky RunEngine.
A subset of the actuators' names must match the names of suggested parameterizations.
sensors: Sequence[Sensor]
Objects that can produce data to acquire data from the beamline using the Bluesky RunEngine.
evaluation_function: EvaluationFunction
A callable that uses an acquisition identifier to retrieve acquired data and produce outcomes.
acquisition_plan: AcquisitionPlan, optional
A Bluesky plan to acquire data from the beamline. If not provided, a default plan will be used.
"""
instance = object.__new__(cls)
instance._optimizer = AxOptimizer.from_checkpoint(checkpoint_path)
instance._actuators = actuators
instance._sensors = sensors
instance._evaluation_function = evaluation_function
instance._acquisition_plan = acquisition_plan
instance._readable_cache = {}
instance._callbacks = [OptimizationLogger()]
instance._callback_router = OptimizationCallbackRouter(instance._callbacks)
return instance
@property
def callbacks(self) -> list[CallbackBase]:
"""The list of active optimization callbacks.
Callbacks in this list receive documents from ``"optimize"`` and
``"sample_suggestions"`` runs. The default list contains an
:class:`~bluesky_adaptive.callbacks.logger.OptimizationLogger`.
The list can be mutated directly, or use :meth:`subscribe` /
:meth:`unsubscribe` for convenience.
"""
return self._callbacks
[docs]
def subscribe(self, callback: CallbackBase) -> None:
"""Subscribe a callback to receive optimization run documents.
Parameters
----------
callback : CallbackBase
A Bluesky callback instance.
Raises
------
ValueError
If *callback* is already subscribed.
"""
if callback in self._callbacks:
raise ValueError(f"Callback {callback!r} is already subscribed.")
self._callbacks.append(callback)
[docs]
def unsubscribe(self, callback: CallbackBase) -> None:
"""Unsubscribe a previously subscribed callback.
Parameters
----------
callback : CallbackBase
The callback instance to remove.
Raises
------
ValueError
If *callback* is not subscribed.
"""
self._callbacks.remove(callback)
@property
def sensors(self) -> Sequence[Sensor]:
"""The sensors used for data acquisition."""
return self._sensors
@property
def actuators(self) -> Sequence[Actuator]:
"""The actuators that control the degrees of freedom."""
return self._actuators
@property
def evaluation_function(self) -> EvaluationFunction[TUid]:
"""The function used to evaluate acquired data and produce outcomes."""
return self._evaluation_function
@property
def acquisition_plan(self) -> AcquisitionPlan[TUid] | None:
"""The acquisition plan for acquiring data, or ``None`` if using the default."""
return self._acquisition_plan
[docs]
def to_optimization_problem(self) -> OptimizationProblem[TUid]:
"""
Construct an optimization problem from the agent.
Creates an immutable :class:`bluesky_adaptive.protocols.OptimizationProblem` that
encapsulates all components needed for optimization. This is typically
used internally by optimization plans.
Returns
-------
OptimizationProblem
An immutable optimization problem that can be deployed via Bluesky.
See Also
--------
bluesky_adaptive.protocols.OptimizationProblem : The optimization problem dataclass.
bluesky_adaptive.plans.optimize : Uses the optimization problem to run optimization.
"""
return OptimizationProblem(
optimizer=self._optimizer,
actuators=self.actuators,
sensors=self.sensors,
evaluation_function=self.evaluation_function,
acquisition_plan=self.acquisition_plan,
)
[docs]
def acquire_baseline(self, parameterization: dict[str, Any] | None = None) -> MsgGenerator[None]:
"""
Acquire a baseline reading for reference.
Acquires data at a specific parameterization (or current positions) to establish
a baseline for comparison. Useful for relative outcome constraints.
Parameters
----------
parameterization : dict[str, Any] | None, optional
The DOF values to move to before acquiring baseline. If None, acquires
at current positions.
Yields
------
Msg
Bluesky messages for the run engine.
See Also
--------
bluesky_adaptive.plans.acquire_baseline : The underlying Bluesky plan.
"""
yield from acquire_baseline(self.to_optimization_problem(), parameterization=parameterization)
[docs]
def optimize(self, iterations: int | None = 1, n_points: int = 1) -> MsgGenerator[None]:
"""
Run Bayesian optimization.
Performs iterative optimization by suggesting points, acquiring data, evaluating
outcomes, and updating the model. This is the main method for running optimization
with an agent.
Parameters
----------
iterations : int | None, optional
The maximum number of optimization iterations to run. If None, run until
the configured stopping strategy's criterion is met. Default is 1. Each
iteration suggests, evaluates, and learns from n_points.
n_points : int, optional
The number of points to evaluate per iteration. Default is 1. Higher values
enable batch optimization but may reduce optimization efficiency per iteration.
Yields
------
Msg
Bluesky messages for the run engine.
Raises
------
ValueError
If ``iterations`` is None and the optimizer does not support stopping criteria.
See Also
--------
bluesky_adaptive.plans.optimize : The underlying Bluesky optimization plan.
suggest : Get point suggestions without running acquisition.
ingest : Manually ingest evaluation results.
Notes
-----
This is the primary method for running optimization. It handles the full loop
of suggesting points, acquiring data, evaluating outcomes, and updating the model.
For complete examples, see :doc:`/tutorials/simple-experiment`.
"""
optimize_plan = optimize(
self.to_optimization_problem(), iterations=iterations, n_points=n_points, readable_cache=self._readable_cache
)
if self._callbacks:
optimize_plan = bpp.subs_wrapper(
optimize_plan,
self._callback_router,
)
yield from optimize_plan
[docs]
def sample_suggestions(
self, suggestions: Sequence[Mapping]
) -> MsgGenerator[tuple[TUid, Sequence[Mapping], Sequence[Mapping]]]:
"""
Evaluate specific parameter combinations.
Acquires data for given suggestions and ingests results. Supports both
optimizer suggestions and manual points.
Parameters
----------
suggestions : Sequence[Mapping]
Either optimizer suggestions (with "_id") or manual points (without "_id").
Returns
-------
tuple[TUid, Sequence[Mapping], Sequence[Mapping]]
Acquisition UID, suggestions with "_id", and outcomes.
See Also
--------
suggest : Get optimizer suggestions.
optimize : Run full optimization loop.
"""
sample_suggestions_plan = sample_suggestions(
self.to_optimization_problem(), suggestions=suggestions, readable_cache=self._readable_cache
)
if self._callbacks:
sample_suggestions_plan = bpp.subs_wrapper(
sample_suggestions_plan,
self._callback_router,
)
return (yield from sample_suggestions_plan)
[docs]
def navigate_to_best(self, parameterization: Mapping | None = None) -> MsgGenerator[None]:
"""
Move actuators to the best point found during optimization.
If no explicit parameterization is provided, queries the optimizer for its
best point(s). For multi-objective optimizers that return multiple Pareto-optimal
points, an explicit parameterization must be provided.
Parameters
----------
parameterization : Mapping | None, optional
Explicit parameterization to navigate to. If None, queries the optimizer's
best point. For multi-objective problems, call ``get_best_points()``
to inspect the Pareto set and select one.
Raises
------
ValueError
If the optimizer returns multiple Pareto-optimal points and no
explicit ``parameterization`` is provided.
See Also
--------
get_best_points : Query the optimizer for best points.
"""
optimization_problem = self.to_optimization_problem()
return (
yield from navigate_to_best(
optimization_problem.actuators,
optimization_problem.optimizer,
parameterization,
)
)