From 754780a3dc0d5cf97f730d4e618fa2e4581cfd07 Mon Sep 17 00:00:00 2001 From: jessica-moylan Date: Mon, 3 Aug 2026 10:05:13 -0400 Subject: [PATCH 1/2] added GlobalStoppingAware protocol for early stopping --- src/blop/plans.py | 10 ++++++++++ src/blop/protocols.py | 20 ++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/blop/plans.py b/src/blop/plans.py index e5910780..3e839dd1 100644 --- a/src/blop/plans.py +++ b/src/blop/plans.py @@ -15,6 +15,7 @@ ID_KEY, Actuator, CanRegisterSuggestions, + GlobalStoppingAware, OptimizationProblem, Sensor, TrialFaultAware, @@ -167,6 +168,7 @@ def optimize( n_points: int = 1, checkpoint_interval: int | None = None, readable_cache: dict[str, InferredReadable] | None = None, + global_stopper: GlobalStoppingAware | None = None, **kwargs: Any, ) -> MsgGenerator[None]: """ @@ -187,6 +189,8 @@ def optimize( readable_cache: dict[str, InferredReadable] | None = None Cache of readable objects to store the suggestions and outcomes as events. If None, a new cache will be created. + global_stopper: GlobalStoppingAware | None = None, + The global stopping strategy to determine when the optimization should stop. **kwargs : Any Additional keyword arguments to pass to the :func:`optimize_step` plan. @@ -218,6 +222,12 @@ def _optimize() -> MsgGenerator[None]: # Perform a single step of the optimization uid, suggestions, outcomes = yield from optimize_step(optimization_problem, n_points, **kwargs) + if isinstance(optimization_problem.optimizer, GlobalStoppingAware): + stop_now, stop_reason = optimization_problem.optimizer.should_stop() + if stop_now: + print(f"Global stopping triggered at iteration {i + 1}: {stop_reason}") + return + # Read the optimization step into the Bluesky and emit events for each suggestion and outcome yield from read_step(uid, suggestions, outcomes, n_points, readable_cache) diff --git a/src/blop/protocols.py b/src/blop/protocols.py index 09d955a5..309be641 100644 --- a/src/blop/protocols.py +++ b/src/blop/protocols.py @@ -81,6 +81,26 @@ def register_failures(self, suggestions: list[dict]) -> None: ... +@runtime_checkable +class GlobalStoppingAware(Protocol): + """ + A protocol for optimizers that can evaluate global 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. + """ + ... + + @runtime_checkable class Checkpointable(Protocol): """ From 0eb411a70005003f69f902532549defc40b5c670 Mon Sep 17 00:00:00 2001 From: jessica-moylan Date: Mon, 3 Aug 2026 14:19:12 -0400 Subject: [PATCH 2/2] make tests and changes based off code review --- src/blop/plans.py | 10 ++-- src/blop/protocols.py | 2 +- src/blop/tests/test_plans.py | 104 +++++++++++++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 7 deletions(-) diff --git a/src/blop/plans.py b/src/blop/plans.py index 3e839dd1..6e39ae61 100644 --- a/src/blop/plans.py +++ b/src/blop/plans.py @@ -15,9 +15,9 @@ ID_KEY, Actuator, CanRegisterSuggestions, - GlobalStoppingAware, OptimizationProblem, Sensor, + StoppingConditions, TrialFaultAware, ) from .utils import InferredReadable, _maybe_checkpoint, collect_optimization_metadata, route_suggestions @@ -168,7 +168,6 @@ def optimize( n_points: int = 1, checkpoint_interval: int | None = None, readable_cache: dict[str, InferredReadable] | None = None, - global_stopper: GlobalStoppingAware | None = None, **kwargs: Any, ) -> MsgGenerator[None]: """ @@ -189,8 +188,6 @@ def optimize( readable_cache: dict[str, InferredReadable] | None = None Cache of readable objects to store the suggestions and outcomes as events. If None, a new cache will be created. - global_stopper: GlobalStoppingAware | None = None, - The global stopping strategy to determine when the optimization should stop. **kwargs : Any Additional keyword arguments to pass to the :func:`optimize_step` plan. @@ -222,10 +219,11 @@ def _optimize() -> MsgGenerator[None]: # Perform a single step of the optimization uid, suggestions, outcomes = yield from optimize_step(optimization_problem, n_points, **kwargs) - if isinstance(optimization_problem.optimizer, GlobalStoppingAware): + if isinstance(optimization_problem.optimizer, StoppingConditions): stop_now, stop_reason = optimization_problem.optimizer.should_stop() if stop_now: - print(f"Global stopping triggered at iteration {i + 1}: {stop_reason}") + reason = stop_reason if stop_reason is not None else "No reason provided" + logger.info(f"Global stopping triggered at iteration {i + 1}: {reason}") return # Read the optimization step into the Bluesky and emit events for each suggestion and outcome diff --git a/src/blop/protocols.py b/src/blop/protocols.py index 309be641..8861a274 100644 --- a/src/blop/protocols.py +++ b/src/blop/protocols.py @@ -82,7 +82,7 @@ def register_failures(self, suggestions: list[dict]) -> None: @runtime_checkable -class GlobalStoppingAware(Protocol): +class StoppingConditions(Protocol): """ A protocol for optimizers that can evaluate global stopping criteria. diff --git a/src/blop/tests/test_plans.py b/src/blop/tests/test_plans.py index 02b190bc..6c49970a 100644 --- a/src/blop/tests/test_plans.py +++ b/src/blop/tests/test_plans.py @@ -11,6 +11,7 @@ EvaluationFunction, OptimizationProblem, Optimizer, + StoppingConditions, TrialFaultAware, ) @@ -477,3 +478,106 @@ def test_acquire_baseline_from_current(RE): optimizer.ingest.assert_called_once_with([{"objective": 0.0, "_id": "baseline", "x1": -1.0}]) assert evaluation_function.call_count == 1 + + +def test_optimize_max_number_of_iterations_before_stop(RE): + """Tests that the optimization stops at a set number of iterations""" + + class StoppingOptimizer(Optimizer, StoppingConditions): ... + + optimizer = MagicMock(spec=StoppingOptimizer) + optimizer.suggest.return_value = [{"x1": 0.0, "_id": 0}] + + # Set up the optimizer to stop after 2 iterations + optimizer.should_stop.side_effect = [ + (False, None), + (True, "converged"), + ] + evaluation_function = MagicMock(spec=EvaluationFunction, return_value=[{"objective": 0.0, "_id": 0}]) + optimization_problem = OptimizationProblem( + optimizer=optimizer, + actuators=[MovableSignal("x1")], + sensors=[ReadableSignal("objective")], + evaluation_function=evaluation_function, + ) + + RE(optimize(optimization_problem, iterations=5)) + + assert optimizer.suggest.call_count == 2 + assert optimizer.should_stop.call_count == 2 + + +def test_optimize_stop_condition_not_hit(RE): + """Tests that optimization stops before stop condition is met""" + + class StoppingOptimizer(Optimizer, StoppingConditions): ... + + optimizer = MagicMock(spec=StoppingOptimizer) + optimizer.suggest.return_value = [{"x1": 0.0, "_id": 0}] + + # Allow for 3 iterations + optimizer.should_stop.side_effect = [(False, None), (False, None), (False, None)] + evaluation_function = MagicMock(spec=EvaluationFunction, return_value=[{"objective": 0.0, "_id": 0}]) + optimization_problem = OptimizationProblem( + optimizer=optimizer, + actuators=[MovableSignal("x1")], + sensors=[ReadableSignal("objective")], + evaluation_function=evaluation_function, + ) + + # We only are running for 2 iterations, so stop condition should not be met + RE(optimize(optimization_problem, iterations=2)) + + assert optimizer.suggest.call_count == 2 + assert optimizer.should_stop.call_count == 2 + + +def test_optimize_stops_when_change_is_within_tolerance(RE): + """Tests that the optimization stops when the change in objective value is within a specified tolerance.""" + + class ToleranceStopOptimizer(Optimizer, StoppingConditions): + def __init__(self, tolerance: float): + self.tolerance = tolerance + self._last_value: float | None = None + self._previous_value: float | None = None + + def suggest(self, num_points: int | None = None) -> list[dict]: + return [{"x1": 0.0, "_id": 0}] + + def ingest(self, points: list[dict]) -> None: + self._previous_value = self._last_value + self._last_value = points[0]["objective"] + + def should_stop(self) -> tuple[bool, str | None]: + if self._previous_value is None or self._last_value is None: + return (False, None) + + if abs(self._last_value - self._previous_value) <= self.tolerance: + return (True, "objective change within tolerance") + + return (False, None) + + # Stop optimization when the change is within 0.1 + optimizer = ToleranceStopOptimizer(tolerance=0.1) + evaluation_function = MagicMock( + spec=EvaluationFunction, + side_effect=[ + [{"objective": 0.5, "_id": 0}], + [{"objective": 0.55, "_id": 0}], + ], + ) + optimization_problem = OptimizationProblem( + optimizer=optimizer, + actuators=[MovableSignal("x1")], + sensors=[ReadableSignal("objective")], + evaluation_function=evaluation_function, + ) + + callback, events = _collect_optimize_events() + RE.subscribe(callback) + try: + RE(optimize(optimization_problem, iterations=5)) + finally: + RE.unsubscribe(callback) + + assert evaluation_function.call_count == 2