Support walltime limit in ScipyOptimizer

SciPy optimizers don't support wall time limits directly.
However, this can be achieved through callbacks. This is done here.

Author: dweindl

pypesto/optimize/optimizer.py

@@ -439,6 +439,8 @@ def __init__(
         if self.options is None:
             self.options = ScipyOptimizer.get_default_options(self)
         self.tol = tol
+        #: maximum walltime in seconds
+        self._maxtime_seconds: float | None = None
 
     def __repr__(self) -> str:
         rep = f"<{self.__class__.__name__} method={self.method}"
@@ -601,6 +603,20 @@ def fun(x):
             if hessp is not None:
                 hess = None
 
+            # Set callback for handling timelimit if necessary
+            callback = None
+            if self._maxtime_seconds is not None and np.isfinite(
+                self._maxtime_seconds
+            ):
+                start_time = time.time()
+
+                def callback(*args, **kwargs):
+                    elapsed_time = time.time() - start_time
+                    if elapsed_time >= self._maxtime_seconds:
+                        raise StopIteration(
+                            f"Maximum time {self._maxtime_seconds}s exceeded."
+                        )
+
             # optimize
             res = scipy.optimize.minimize(
                 fun=fun,
@@ -612,6 +628,7 @@ def fun(x):
                 bounds=bounds,
                 options=self.options,
                 tol=self.tol,
+                callback=callback,
             )
             # extract fval/grad from result
             grad = getattr(res, "jac", None)
@@ -670,6 +687,41 @@ def set_maxiter(self, iterations: int) -> None:
         else:
             self.options["maxiter"] = iterations
 
+    def supports_maxtime(self) -> bool:
+        """
+        Check whether optimizer supports time limits.
+
+        Returns
+        -------
+        True if optimizer supports setting a maximum wall time,
+        False otherwise.
+        """
+        # TNC neither supports time limits nor callback functions
+        return self.method.lower() != "tnc"
+
+    def set_maxtime(self, seconds: float) -> None:
+        """
+        Set the maximum wall time for optimization.
+
+        Parameters
+        ----------
+        seconds
+            Maximum wall time in seconds.
+
+        Raises
+        ------
+        NotImplementedError
+            If the optimizer does not support time limits.
+        """
+        if not self.supports_maxtime():
+            raise NotImplementedError(
+                f"{self.__class__.__name__} method {self.method} does not "
+                "support time limits. "
+                f"Check supports_maxtime() before calling set_maxtime()."
+            )
+
+        self._maxtime_seconds = seconds
+
 
 class IpoptOptimizer(Optimizer):
     """Use Ipopt (https://pypi.org/project/cyipopt/) for optimization."""

test/optimize/test_optimizer_common_interface.py

@@ -1,21 +1,16 @@
 """Test unified interface for optimizer limits (time, iterations, evaluations)."""
 
+import numpy as np
 import pytest
+import scipy.optimize as so
 
+import pypesto
 import pypesto.optimize as optimize
 
 
 class TestOptimizerMaxtimeInterface:
     """Test the unified maxtime interface for optimizers."""
 
-    def test_scipy_optimizer_no_support(self):
-        """Test that ScipyOptimizer does not support time limits."""
-        optimizer = optimize.ScipyOptimizer()
-        assert optimizer.supports_maxtime() is False
-
-        with pytest.raises(NotImplementedError):
-            optimizer.set_maxtime(10.0)
-
     def test_ipopt_optimizer_support(self):
         """Test IpoptOptimizer time limit support."""
         optimizer = optimize.IpoptOptimizer()
@@ -71,6 +66,32 @@ def test_ess_optimizer_support(self):
         optimizer.set_maxtime(10.0)
         assert optimizer.max_walltime_s == 10.0
 
+    def test_scipy_optimizer_support(self):
+        """Test ScipyOptimizer time limit support."""
+        optimizer = optimize.ScipyOptimizer()
+        assert optimizer.supports_maxtime() is True
+        optimizer.set_maxtime(1.0)
+
+        def fun(x):
+            import time
+
+            time.sleep(0.1)
+            return so.rosen(x)
+
+        objective = pypesto.Objective(fun=fun, grad=so.rosen_der)
+        dim_full = 2
+        lb = -5 * np.ones((dim_full, 1))
+        ub = 5 * np.ones((dim_full, 1))
+
+        problem = pypesto.Problem(
+            objective=objective,
+            lb=lb,
+            ub=ub,
+        )
+
+        result = optimizer.minimize(problem, id="test", x0=np.zeros(dim_full))
+        assert "StopIteration" in result.message
+
 
 class TestOptimizerMaxiterInterface:
     """Test the unified maxiter interface for optimizers."""