diff --git a/src/blop/plans.py b/src/blop/plans.py index e5910780..6e39ae61 100644 --- a/src/blop/plans.py +++ b/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) diff --git a/src/blop/protocols.py b/src/blop/protocols.py index 09d955a5..8861a274 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 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): """ 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