OptimizationProblem dataclass and updated plans

Author: thopkins32

src/blop/ax/agent.py

@@ -10,8 +10,8 @@
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from bluesky.protocols import Readable
 
-from ..evaluation import default_evaluation_function
 from ..dofs import DOF, DOFConstraint
+from ..evaluation import default_evaluation_function
 from ..objectives import Objective
 from .adapters import configure_metrics, configure_objectives, configure_parameters
 

src/blop/evaluation.py

@@ -1,8 +1,8 @@
 from ax.api.types import TOutcome
 from tiled.client.container import Container
 
-from .objectives import Objective
 from .data_access import TiledDataAccess
+from .objectives import Objective
 
 
 def default_evaluation_function(

src/blop/plans.py

@@ -1,55 +1,46 @@
 import functools
-from collections import defaultdict
+import warnings
 from collections.abc import Callable, Generator, Mapping, Sequence
-from typing import Any
+from typing import Any, cast
 
 import bluesky.plan_stubs as bps
 import bluesky.plans as bp
-from ax.api.types import TParameterization, TParameterValue
-from bluesky.protocols import Movable, Readable, Reading
-from bluesky.run_engine import Msg
-from bluesky.utils import MsgGenerator, plan
+from ax.api.types import TParameterization
+from bluesky.protocols import Movable, NamedMovable, Readable, Reading
+from bluesky.utils import Msg, MsgGenerator, plan
 from ophyd import Signal  # type: ignore[import-untyped]
 
 from .dofs import DOF
-from .ax.agent import Agent
+from .protocols import OptimizationProblem
 
 
-def _unpack_parameters(dofs: dict[str, DOF], parameterizations: list[TParameterization]) -> list[Movable | TParameterValue]:
-    """Unpack the parameterizations into Bluesky plan arguments."""
-    unpacked_dict = defaultdict(list)
-    for parameterization in parameterizations:
-        for dof_name in dofs.keys():
-            if dof_name in parameterization:
-                unpacked_dict[dof_name].append(parameterization[dof_name])
-            else:
-                raise ValueError(f"Parameter {dof_name} not found in parameterization. Parameterization: {parameterization}")
-
+def _unpack_for_list_scan(movables: Mapping[NamedMovable, Sequence[Any]]) -> list[NamedMovable | Any]:
+    """Unpack the movables and inputs into Bluesky list_scan plan arguments."""
     unpacked_list = []
-    for dof_name, values in unpacked_dict.items():
-        unpacked_list.append(dofs[dof_name].movable)
+    for movable, values in movables.items():
+        unpacked_list.append(movable)
         unpacked_list.append(values)
 
     return unpacked_list
 
 
 @plan
-def acquire(
+def default_acquire(
+    movables: Mapping[NamedMovable, Sequence[Any]],
     readables: Sequence[Readable],
-    dofs: dict[str, DOF],
-    trials: dict[int, TParameterization],
+    *,
     per_step: bp.PerStep | None = None,
     **kwargs: Any,
 ) -> MsgGenerator[str]:
     """
-    A plan to acquire data for optimization.
+    A default plan to acquire data for optimization. Simply a list scan.
 
     Parameters
     ----------
+    movables: Mapping[NamedMovable, Sequence[Any]]
+        The movables to move and the inputs to move them to.
     readables: Sequence[Readable]
         The readables to trigger and read.
-    dofs: dict[str, DOF]
-        A dictionary mapping DOF names to DOFs.
     trials: dict[int, TParameterization]
         A dictionary mapping trial indices to their suggested parameterizations. Typically only a single trial is provided.
     per_step: bp.PerStep | None = None
@@ -66,69 +57,81 @@ def acquire(
     --------
     bluesky.plans.list_scan : The Bluesky plan to acquire data.
     """
-    plan_args = _unpack_parameters(dofs, trials.values())
+    plan_args = _unpack_for_list_scan(movables)
     return (
+        # TODO: fix argument type in bluesky.plans.list_scan
         yield from bp.list_scan(
-            readables, *plan_args, md={"ax_trial_indices": list(trials.keys())}, per_step=per_step, **kwargs
+            readables,
+            *plan_args,
+            per_step=per_step,
+            **kwargs,  # type: ignore[arg-type]
         )
     )
 
 
 @plan
 def optimize_step(
-    generator: Agent,
-    n: int = 1,
-    acquisition_plan: Callable[[], MsgGenerator[None]] | None = None,
+    optimization_problem: OptimizationProblem,
+    n_points: int = 1,
+    *args: Any,
+    **kwargs: Any,
 ) -> MsgGenerator[None]:
     """
     A single step of the optimization loop.
 
     Parameters
     ----------
-    generator : Agent
-        The generator to optimize with.
-    n : int, optional
-        The number of trials to suggest.
-    acquisition_plan : Callable[[], MsgGenerator[None]] | None, optional
-        The acquisition plan to use to acquire data. If not provided, the default acquisition plan will be used.
+    optimization_problem : OptimizationProblem
+        The optimization problem to solve.
+    n_points : int, optional
+        The number of points to suggest.
     """
-    if acquisition_plan is None:
-        acquisition_plan = acquire
-    trials = generator.suggest(n)
-    data = yield from acquisition_plan(generator.readables, generator.dofs, trials)
-    outcomes = generator.evaluate(trials, data)
+    if optimization_problem.acquisition_plan is None:
+        acquisition_plan = default_acquire
+    else:
+        acquisition_plan = optimization_problem.acquisition_plan
+    generator = optimization_problem.generator
+    movables = optimization_problem.movables
+    suggestions = generator.suggest(n_points)
+    movables_and_inputs = {movable: [suggestion[movable.name] for suggestion in suggestions] for movable in movables}
+    uid = yield from acquisition_plan(movables_and_inputs, optimization_problem.readables, *args, **kwargs)
+    outcomes = optimization_problem.evaluation_function(uid)
     generator.ingest(outcomes)
 
 
 @plan
 def optimize(
-    generator: Agent,
+    optimization_problem: OptimizationProblem,
     iterations: int = 1,
-    n: int = 1,
-    acquisition_plan: Callable[[], MsgGenerator[None]] | None = None,
+    n_points: int = 1,
+    *args: Any,
+    **kwargs: Any,
 ) -> MsgGenerator[None]:
     """
     A plan to optimize the generator.
 
     Parameters
     ----------
-    generator : Agent
-        The generator to optimize with.
+    optimization_problem : OptimizationProblem
+        The optimization problem to solve.
     iterations : int, optional
         The number of optimization iterations to run.
-    n : int, optional
-        The number of trials to suggest per iteration.
-    acquisition_plan : Callable[[], MsgGenerator[None]] | None, optional
-        The acquisition plan to use to acquire data. If not provided, the default acquisition plan will be used.
+    n_points : int, optional
+        The number of points to suggest per iteration.
     """
 
     for _ in range(iterations):
-        yield from optimize_step(generator, n, acquisition_plan)
+        yield from optimize_step(optimization_problem, n_points, *args, **kwargs)
 
 
 @plan
 def list_scan_with_delay(*args: Any, delay: float = 0, **kwargs: Any) -> Generator[Msg, None, str]:
-    "Accepts all the normal 'scan' parameters, plus an optional delay."
+    """
+    Accepts all the normal 'scan' parameters, plus an optional delay.
+
+    .. deprecated:: v0.8.2
+        This plan is deprecated and will be removed in Blop v1.0.0. See documentation how-to-guides for more information.
+    """
 
     def one_nd_step_with_delay(
         detectors: Sequence[Signal], step: Mapping[Movable, Any], pos_cache: Mapping[Movable, Any]
@@ -149,6 +152,11 @@ def default_acquisition_plan(
     dofs: Sequence[DOF], inputs: Mapping[str, Sequence[Any]], dets: Sequence[Signal], **kwargs: Any
 ) -> Generator[Msg, None, str]:
     """
+    Default acquisition plan.
+
+    .. deprecated:: v0.8.2
+        This plan is deprecated and will be removed in Blop v1.0.0. See documentation how-to-guides for more information.
+
     Parameters
     ----------
     x : list of DOFs or DOFList
@@ -158,6 +166,11 @@ def default_acquisition_plan(
     dets: list
         A list of detectors to trigger
     """
+    warnings.warn(
+        "This plan is deprecated and will be removed in Blop v1.0.0. See documentation how-to-guides for more information.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     delay = kwargs.get("delay", 0)
     args = []
     for dof in dofs:
@@ -177,6 +190,11 @@ def read(readables: Sequence[Readable], **kwargs: Any) -> MsgGenerator[dict[str,
     ----------
     readables : Sequence[Readable]
         The readables to read.
+
+    Returns
+    -------
+    dict[str, Any]
+        A dictionary of the readable names and their current values.
     """
     results = {}
     for readable in readables:
@@ -281,18 +299,16 @@ def acquire_with_background(
     per_step_background_read : The per-step plan to take background readings.
     """
     per_step = per_step_background_read(block_beam, unblock_beam)
-    return (yield from acquire(readables, dofs, trials, per_step=per_step, **kwargs))
+    return (yield from default_acquire(readables, dofs, trials, per_step=per_step, **kwargs))
 
 
 def acquire_baseline(
-    generator: Agent,
-    parameterization: TParameterization | None = None,
-    arm_name: str | None = None,
-    acquisition_plan: Callable[[Sequence[Readable], dict[str, DOF], dict[int, TParameterization], Any], MsgGenerator[str]] | None = None,
+    optimization_problem: OptimizationProblem,
+    parameterization: dict[str, Any] | None = None,
     **kwargs: Any,
 ) -> MsgGenerator[None]:
     """
-    Acquire a baseline reading.
+    Acquire a baseline reading. Useful for relative outcome constraints.
 
     Parameters
     ----------
@@ -306,10 +322,26 @@ def acquire_baseline(
         The per-step plan to execute for each step of the scan.
     **kwargs: Any
         Additional keyword arguments to pass to the acquire plan.
+
+    See Also
+    --------
+    default_acquire : The default plan to acquire data.
     """
+    movables = optimization_problem.movables
     if parameterization is None:
-        parameterization = yield from read([dof.movable for dof in generator.dofs.values()])
-    trials = generator.attach_baseline(parameters=parameterization, arm_name=arm_name)
-    uid = yield from acquisition_plan(generator.readables, generator.dofs, trials, **kwargs)
-    outcomes = generator.evaluation_function(trials, uid, **generator.evaluation_kwargs)
-    generator.ingest(trials, outcomes)
+        if all(isinstance(movable, Readable) for movable in movables):
+            parameterization = yield from read(cast(Sequence[Readable], movables))
+        else:
+            raise ValueError(
+                "All movables must also implement the Readable protocol to acquire a baseline from current positions."
+            )
+    generator = optimization_problem.generator
+    if optimization_problem.acquisition_plan is None:
+        acquisition_plan = default_acquire
+    else:
+        acquisition_plan = optimization_problem.acquisition_plan
+    movables_and_inputs = {movable: parameterization[movable.name] for movable in movables}
+    uid = yield from acquisition_plan(movables_and_inputs, optimization_problem.readables, **kwargs)
+    outcome = optimization_problem.evaluation_function(uid)[0]
+    outcome.update({"_id": 0})
+    generator.ingest([outcome])

src/blop/protocols.py

@@ -1,11 +1,16 @@
-from collections.abc import Sequence
-from typing import Any, Mapping, Protocol
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
 
 from bluesky.protocols import NamedMovable, Readable
 from bluesky.utils import MsgGenerator, plan
 
 
-class Agent(Protocol):
+@runtime_checkable
+class Generator(Protocol):
+    """
+    A minimal generator interface for optimization.
+    """
 
     def suggest(self, num_points: int | None = None) -> list[dict]:
         """
@@ -38,6 +43,7 @@ def ingest(self, points: list[dict]) -> None:
         ...
 
 
+@runtime_checkable
 class EvaluationFunction(Protocol):
     def __call__(self, uid: str, *args: Any, **kwargs: Any) -> list[dict]:
         """
@@ -56,6 +62,7 @@ def __call__(self, uid: str, *args: Any, **kwargs: Any) -> list[dict]:
         ...
 
 
+@runtime_checkable
 class AcquisitionPlan(Protocol):
     @plan
     def __call__(
@@ -67,7 +74,7 @@ def __call__(
     ) -> MsgGenerator[str]:
         """
         Acquire data for optimization.
-        
+
         This should be a Bluesky plan that moves the movables to each of their suggested positions
         and acquires data from the readables.
 
@@ -83,4 +90,31 @@ def __call__(
         str
             The unique identifier of the Bluesky run.
         """
-        ...
+        ...
+
+
+@dataclass(frozen=True)
+class OptimizationProblem:
+    """
+    An optimization problem to solve. Immutable once initialized.
+
+    Attributes
+    ----------
+    generator: Generator
+        Suggests points to evaluate and ingests outcomes to inform the optimization.
+    movables: Sequence[NamedMovable]
+        Objects that can be moved to control the beamline using the Bluesky RunEngine.
+        A subset of the movables' names must match the names of suggested parameterizations.
+    readables: Sequence[Readable]
+        Objects that can be read to acquire data from the beamline using the Bluesky RunEngine.
+    evaluation_function: EvaluationFunction
+        A callable to evaluate data from a Bluesky run 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.
+    """
+
+    generator: Generator
+    movables: Sequence[NamedMovable]
+    readables: Sequence[Readable]
+    evaluation_function: EvaluationFunction
+    acquisition_plan: AcquisitionPlan | None = None