add GlobalStoppingAware protocol

src/blop/plans.py

@@ -17,6 +17,7 @@
     CanRegisterSuggestions,
     OptimizationProblem,
     Sensor,
+    StoppingConditions,
     TrialFaultAware,
 )
 from .utils import InferredReadable, _maybe_checkpoint, collect_optimization_metadata, route_suggestions
@@ -218,6 +219,13 @@ 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, StoppingConditions):
+                stop_now, stop_reason = optimization_problem.optimizer.should_stop()
+                if stop_now:
+                    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
             yield from read_step(uid, suggestions, outcomes, n_points, readable_cache)
 

src/blop/protocols.py

@@ -81,6 +81,26 @@ def register_failures(self, suggestions: list[dict]) -> None:
         ...
 
 
+@runtime_checkable
+class StoppingConditions(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):
     """

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