"""Protocols that bridge between optimizer backends and Bluesky."""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Generic, Literal, Protocol, TypeVar, runtime_checkable
from bluesky.protocols import EventCollectable, EventPageCollectable, Flyable, HasName, Movable, Readable
from bluesky.utils import MsgGenerator, plan
[docs]
@runtime_checkable
class MovableHasName(Movable, HasName, Protocol):
"""
A movable that has a name.
We use this instead of `bluesky.protocols.NamedMovable` since
we do not want to require `HasHints` on the movable.
A `Movable` and `HasName` is sufficient. `HasHints` should be optional.
"""
...
ID_KEY: Literal["_id"] = "_id"
Actuator = MovableHasName | Flyable
Sensor = Readable | EventCollectable | EventPageCollectable
TActuator = TypeVar("TActuator")
"""Actuator generic type"""
TSensor = TypeVar("TSensor")
"""Sensor generic type"""
TPlan = TypeVar("TPlan")
"""Plan generic type"""
TUid = TypeVar("TUid")
"""Acquisition UID generic type"""
_TUid_contra = TypeVar("_TUid_contra", contravariant=True)
"""UID TypeVar for EvaluationFunction variance."""
_TUid_co = TypeVar("_TUid_co", covariant=True)
"""UID TypeVar for AcquisitionPlan variance."""
[docs]
@runtime_checkable
class CanRegisterSuggestions(Protocol):
"""
A protocol for optimizers that can register suggestions.
This allows them to add an "_id" key to the suggestions dynamically and ensure
that the suggestions are unique.
"""
[docs]
def register_suggestions(self, suggestions: Sequence[Mapping]) -> Sequence[Mapping]:
"""
Register the suggestions with the optimizer.
Parameters
----------
suggestions: Sequence[Mapping]
The suggestions to register. The "_id" key is optional and will be overwritten if present.
Returns
-------
Sequence[Mapping]
The suggestions with an "_id" key added.
"""
...
[docs]
@runtime_checkable
class TrialFaultAware(Protocol):
"""
A protocol to accept information about trial failures of the optimization loop.
Used to invalidate or register early stop on data for the optimizer, or do necesary
cleanup of processes not directly tied to the run engine
"""
[docs]
def register_failures(self, suggestions: Sequence[Mapping]) -> None:
"""
Register the failed suggestions with the optimizer.
Parameters
----------
suggestions: Sequence[Mapping]
The suggestions to fail. Ids must be present.
"""
...
@runtime_checkable
class SupportsStoppingCriteria(Protocol):
"""
A protocol for optimizers that can evaluate stopping criteria.
This allows optimization plans to terminate early through tolerance, max_iterations etc.
"""
def should_stop(self) -> tuple[bool, str | None]:
"""
Evaluate whether optimization should terminate early.
Returns
-------
tuple[bool, str | None]
(stop_now, reason). If stop_now is True, optimization should stop. If reason is provided, it will be logged.
"""
...
[docs]
@runtime_checkable
class Checkpointable(Protocol):
"""
A protocol for objects that can can write state to persistent storage.
Implementers configure storage at construction time (e.g., a file path, databse URI).
The checkpoint method then saves or updates to that pre-configured location.
"""
[docs]
def checkpoint(self) -> None:
"""Write the object's state to persistent storage."""
...
[docs]
@runtime_checkable
class Optimizer(Protocol):
"""
A minimal optimizer interface for optimization.
This protocol defines the interface for optimizers in Bluesky Adaptive. Most users will use
the built-in :class:`bluesky_adaptive.ax.optimizer.AxOptimizer`, which provides Bayesian optimization
via Ax. Custom implementations are only needed for specialized optimization algorithms.
See Also
--------
bluesky_adaptive.ax.optimizer.AxOptimizer : Built-in Ax-based optimizer implementation.
bluesky_adaptive.ax.Agent : High-level interface that uses AxOptimizer internally.
"""
[docs]
def suggest(self, num_points: int | None = None) -> Sequence[Mapping]:
"""
Suggest a set of points in the input space, to be evaulated next.
The "_id" key is optional and can be used to identify suggested trials for later evaluation
and ingestion. When optimization plans require IDs, each "_id" value must be unique within
the batch and hashable.
Parameters
----------
num_points : int | None, optional
The number of points to suggest. If not provided, will default to 1.
Returns
-------
Sequence[Mapping]
A sequence of mappings, each containing a parameterization of a point to evaluate next.
Each mapping must contain a unique, hashable "_id" key to identify each parameterization.
"""
...
[docs]
def ingest(self, points: Sequence[Mapping]) -> None:
"""
Ingest a set of points into the experiment. Either from previously suggested points or from an external source.
The "_id" key is optional and can be used to identify points from previously suggested trials or to identify
the point as a "baseline" trial.
Parameters
----------
points : Sequence[Mapping]
A sequence of mappings containing the outcomes of the evaluated parameterizations.
"""
...
[docs]
def get_best_points(self) -> Sequence[tuple[Any, Mapping, Mapping]]:
"""
Get a sequence 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[Any, Mapping, Mapping]]
Each element is a tuple of:
- "_id" of the suggestion
- suggested parameters
- measured outcomes
"""
...
[docs]
@runtime_checkable
class EvaluationFunction(Protocol[_TUid_contra]):
"""
A protocol for transforming acquired data into measurable outcomes.
This protocol defines how to extract and compute optimization objectives from
acquired data. Custom implementations are needed to define how your beamline
data translates into the outcomes you want to optimize.
See Also
--------
bluesky_adaptive.ax.Agent : Accepts an evaluation function during initialization.
Notes
-----
The evaluation function is called after data acquisition to compute outcomes.
It uses the uid returned by the acquisition plan to retrieve or identify
the relevant data and associate each acquired row with a suggestion by its
``"_id"``. The uid type is evaluator-defined; it may be a Bluesky run UID,
ordered suggestion IDs, event UIDs, or a backend-specific UID type.
Examples
--------
See the tutorial documentation for complete examples of evaluation functions:
:doc:`/tutorials/simple-experiment`
"""
[docs]
def __call__(self, uid: _TUid_contra, suggestions: Sequence[Mapping]) -> Sequence[Mapping]:
"""
Evaluate the acquired data and produce outcomes.
Parameters
----------
uid: _TUid_contra
The acquisition UID returned by the acquisition plan.
suggestions: Sequence[Mapping]
A sequence of mappings, each containing a parameterization to evaluate.
Each mapping must contain a unique ``"_id"``. Do not assume its order
aligns with acquired data or preserves optimizer generation order.
Returns
-------
Sequence[Mapping]
A sequence of mappings containing the outcomes of the acquisition, one for each
suggested parameterization. Each mapping must contain an ``"_id"`` identifying
its evaluated suggestion.
"""
...
[docs]
@runtime_checkable
class AcquisitionPlan(Protocol[_TUid_co]):
"""
A protocol for custom data acquisition plans.
This protocol defines how to acquire data from the beamline. Most users will use
the default :func:`bluesky_adaptive.plans.default_acquire` plan, which performs a list scan
in its own Bluesky run, or :func:`bluesky_adaptive.plan_stubs.list_scan_in_run`, which performs
a list scan inside an already-open optimization run. Custom implementations are only
needed for specialized acquisition strategies (e.g., fly scans, complex detector configurations).
See Also
--------
bluesky_adaptive.plans.default_acquire : Default run-owning acquisition plan implementation.
bluesky_adaptive.plan_stubs.list_scan_in_run : Default in-run acquisition plan implementation.
bluesky_adaptive.ax.Agent : Accepts an optional acquisition plan during initialization.
Notes
-----
The acquisition plan is a Bluesky plan that should move the actuators to each
suggested position, acquire data from the sensors, and return a uid that the
evaluation function can use to retrieve or identify the acquired data. When
it may reorder points, it must record their IDs in actual acquisition order.
"""
[docs]
@plan
def __call__(
self,
suggestions: Sequence[Mapping],
actuators: Sequence[Actuator],
sensors: Sequence[Sensor] | None = None,
md: Mapping[str, Any] | None = None,
) -> MsgGenerator[_TUid_co]:
"""
Acquire data for optimization.
This should be a Bluesky plan that moves the actuators to each of their suggested positions
and acquires data from the sensors. Suggestions may be re-ordered for more efficient acquisition,
but it is the responsibility of the implementer to return a uid that lets the matching
evaluation function correlate acquired data with suggestion IDs.
Parameters
----------
suggestions: Sequence[Mapping]
A sequence of mappings, each containing the parameterization of a point to evaluate.
The "_id" key is optional and can be used to identify each suggestion. When present, "_id"
values used by Bluesky Adaptive optimization plans must be unique within a batch and hashable.
actuators: Sequence[Actuator]
The actuators to move to their suggested positions.
sensors: Sequence[Sensor], optional
The sensors that produce data to evaluate.
md : Mapping[str, Any] | None, optional
Metadata to attach to the start document
Returns
-------
same UID type
The uid passed unchanged to the evaluation function. Examples include a Bluesky run UID,
ordered suggestion IDs, event UIDs, or a backend-specific acquisition object.
"""
...
[docs]
@dataclass(frozen=True)
class BaseOptimizationProblem(Generic[TActuator, TSensor, TUid, TPlan]):
"""Base class for optimization problem definitions.
Provides the common structure shared by all optimization problem types.
Users should use the concrete subclasses :class:`OptimizationProblem` or
:class:`QueueserverOptimizationProblem` instead of this class directly.
See Also
--------
OptimizationProblem : Concrete problem type for standard usage.
QueueserverOptimizationProblem : Concrete problem type for queue server usage.
"""
optimizer: Optimizer
actuators: Sequence[TActuator]
sensors: Sequence[TSensor]
evaluation_function: EvaluationFunction[TUid]
acquisition_plan: TPlan | None = None
[docs]
@dataclass(frozen=True)
class OptimizationProblem(BaseOptimizationProblem[Actuator, Sensor, TUid, AcquisitionPlan[TUid]], Generic[TUid]):
"""
An optimization problem to solve. Immutable once initialized.
This dataclass encapsulates all components needed for optimization into a single
immutable structure. It is typically created via :meth:`bluesky_adaptive.ax.Agent.to_optimization_problem`
and used with optimization plans like :func:`bluesky_adaptive.plans.optimize`.
Attributes
----------
optimizer: Optimizer
Suggests points to evaluate and ingests outcomes to inform the optimization.
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[TUid]
A callable that uses an acquisition UID to retrieve acquired data and produce outcomes.
acquisition_plan: AcquisitionPlan[TUid], optional
A Bluesky plan to acquire data from the beamline. If not provided, a default plan will be used.
See Also
--------
bluesky_adaptive.ax.Agent.to_optimization_problem : Creates an OptimizationProblem from an Agent.
bluesky_adaptive.plans.optimize : Bluesky plan that uses an OptimizationProblem.
"""
...
[docs]
@dataclass(frozen=True)
class QueueserverOptimizationProblem(BaseOptimizationProblem[str, str, str, str]):
"""
An optimization problem to solve. Immutable once initialized.
This dataclass encapsulates all components needed for optimization into a single
immutable structure. It is typically created via
:meth:`bluesky_adaptive.ax.queueserver_agent.QueueserverAgent.to_optimization_problem`
and used with bluesky-queueserver-api. Actuators, sensors, and the acquisition plan are referenced
by their names, since their instances live on a remote server.
Attributes
----------
optimizer: Optimizer
Suggests points to evaluate and ingests outcomes to inform the optimization.
actuators: Sequence[str]
Names of 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[str]
Names of objects that can produce data to acquire data from the beamline using the Bluesky RunEngine.
evaluation_function: EvaluationFunction[str]
A callable that uses a Bluesky run UID to retrieve acquired data and produce outcomes.
acquisition_plan: str, optional
The name of a Bluesky plan to acquire data. If not provided, a default plan name will be used.
The plan must match the arguments of :class:`AcquisitionPlan`.
acquisition_plan_kwargs: Mapping[str, Any], optional
Additional plan arguments to pass to the Bluesky plan.
See Also
--------
bluesky_adaptive.ax.queueserver_agent.QueueserverAgent.to_optimization_problem :
Creates a QueueserverOptimizationProblem from an agent.
bluesky_adaptive.queueserver.QueueserverOptimizationRunner :
Runs the optimization loop using the bluesky-queueserver-api.
"""
acquisition_plan_kwargs: Mapping[str, Any] | None = None