Add time limit and solver status

chrhansk · Nov 9, 2023 · 5ee3585 · 5ee3585
1 parent 93d57a9
commit 5ee3585
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 5 deletions.
diff --git a/pygradflow/params.py b/pygradflow/params.py
@@ -87,6 +87,8 @@ class Params:
 
     validate_input: bool = True
 
+    time_limit: float = np.inf
+
     @property
     def dtype(self):
         return np.float32 if self.precision == Precision.Single else np.float64
diff --git a/pygradflow/solver.py b/pygradflow/solver.py
@@ -1,4 +1,7 @@
 from collections import namedtuple
+import time
+
+from enum import Enum, auto
 
 import numpy as np
 from termcolor import colored
@@ -14,14 +17,27 @@
     DualNormUpdate,
     PenaltyStrategy,
 )
+
 from pygradflow.step.step_control import (
     StepResult,
     step_controller,
     StepController,
 )
 
 
-Result = namedtuple("Result", ["x", "y", "d", "success"])
+class SolverStatus(Enum):
+    Converged = auto(),
+    IterationLimit = auto(),
+    TimeLimit = auto()
+
+
+class Result:
+    def __init__(self, x, y, d, success, status):
+        self.x = x
+        self.y = y
+        self.d = d
+        self.success = success
+        self.status = status
 
 
 def bold(s: str) -> str:
@@ -155,12 +171,21 @@ def solve(self, x_0: np.ndarray, y_0: np.ndarray) -> Result:
 
         logger.info("Initial Aug Lag: %.10e", iterate.aug_lag(self.rho))
 
+        status = None
+        start_time = time.time()
+
         for i in range(params.num_it):
             if (i % 25) == 0:
                 print_header()
 
             if iterate.total_res <= params.opt_tol:
                 logger.info("Convergence achieved")
+                status = SolverStatus.Converged
+                break
+
+            if time.time() - start_time >= params.time_limit:
+                logger.info("Reached time limit")
+                status = SolverStatus.TimeLimit
                 break
 
             step_result = self.compute_step(controller, iterate, 1.0 / lamb)
@@ -201,9 +226,9 @@ def solve(self, x_0: np.ndarray, y_0: np.ndarray) -> Result:
                 next_rho = self.penalty.update(iterate, next_iterate)
 
                 if next_rho != self.rho:
-                    logger.debug(
-                        "Updating penalty parameter from %e to %e", self.rho, next_rho
-                    )
+                    logger.debug("Updating penalty parameter from %e to %e",
+                                 self.rho,
+                                 next_rho)
                     self.rho = next_rho
 
                 delta = iterate.dist(next_iterate)
@@ -212,10 +237,12 @@ def solve(self, x_0: np.ndarray, y_0: np.ndarray) -> Result:
 
                 if (lamb <= params.lamb_term) and (delta <= params.opt_tol):
                     logger.info("Convergence achieved")
+                    status = SolverStatus.Converged
                     break
 
         else:
             success = False
+            status = SolverStatus.IterationLimit
             logger.info("Iteration limit reached")
 
         self.print_result(iterate)
@@ -224,4 +251,6 @@ def solve(self, x_0: np.ndarray, y_0: np.ndarray) -> Result:
         y = iterate.y
         d = iterate.bound_duals
 
-        return Result(x, y, d, success)
+        assert status is not None
+
+        return Result(x, y, d, success, status)