Clean Up Newton Solver

wntr/sim/aml/aml.py

@@ -320,7 +320,7 @@ def _remove_constraint(self, con):
         del self._params_referenced_by_con[con]
         del self._floats_referenced_by_con[con]
 
-    def evaluate_residuals(self, x=None, num_threads=4):
+    def evaluate_residuals(self, x=None):
         if x is not None:
             self._evaluator.load_var_values_from_x(x)
         r = self._evaluator.evaluate(len(self._con_ccon_map))

wntr/sim/core.py

@@ -1194,19 +1194,11 @@ def run_sim(self, solver=NewtonSolver, backup_solver=None, solver_options=None,
         Parameters
         ----------
         solver: object
-            wntr.sim.solvers.NewtonSolver or Scipy solver
+            :py:class:`~wntr.sim.solvers.NewtonSolver` or Scipy solver
         backup_solver: object
-            wntr.sim.solvers.NewtonSolver or Scipy solver
+            :py:class:`~wntr.sim.solvers.NewtonSolver` or Scipy solver
         solver_options: dict
-            Solver options are specified using the following dictionary keys:
-
-            * MAXITER: the maximum number of iterations for each hydraulic solve (each timestep and trial) (default = 3000)
-            * TOL: tolerance for the hydraulic equations (default = 1e-6)
-            * BT_RHO: the fraction by which the step length is reduced at each iteration of the line search (default = 0.5)
-            * BT_MAXITER: the maximum number of iterations for each line search (default = 100)
-            * BACKTRACKING: whether or not to use a line search (default = True)
-            * BT_START_ITER: the newton iteration at which a line search should start being used (default = 2)
-            * THREADS: the number of threads to use in constraint and jacobian computations
+            See :py:class:`~wntr.sim.solvers.NewtonSolver` for possible options
         backup_solver_options: dict
         convergence_error: bool (optional)
             If convergence_error is True, an error will be raised if the

wntr/sim/solvers.py

@@ -3,8 +3,11 @@
 import warnings
 import logging
 import enum
+import time
 
-warnings.filterwarnings("error",'Matrix is exactly singular', sp.linalg.MatrixRankWarning)
+warnings.filterwarnings(
+    "error", "Matrix is exactly singular", sp.linalg.MatrixRankWarning
+)
 np.set_printoptions(precision=3, threshold=10000, linewidth=300)
 
 logger = logging.getLogger(__name__)
@@ -18,49 +21,99 @@ class SolverStatus(enum.IntEnum):
 class NewtonSolver(object):
     """
     Newton Solver class.
+
+    Attributes
+    ----------
+    log_progress: bool
+        If True, the infinity norm of the constraint violation will be logged each iteration
+    log_level: int
+        The level for logging the infinity norm of the constraint violation
+    time_limit: float
+        If the wallclock time exceeds time_limit, the newton solver will exit with an error status
+    maxiter: int
+        If the number of iterations exceeds maxiter, the newton solver will exit with an error status
+    tol: float
+        The convergence tolerance. If the infinity norm of the constraint violation drops below tol,
+        the newton solver will exit with a converged status.
+    rho: float
+        During the line search, rho is used to reduce the stepsize. It should be strictly between 0 and 1.
+    bt_maxiter: int
+        The maximum number of line search iterations for each outer iteration
+    bt: bool
+        If False, a line search will not be used.
+    bt_start_iter: int
+        A line search will not be used for any iteration prior to bt_start_iter
     """
-    
+
     def __init__(self, options=None):
+        """
+        Parameters
+        ----------
+        options: dict
+            A dictionary specifying options for the newton solver. Keys
+            should be strings in all caps. See the documentation of the
+            NewtonSolver attributes for details on each option. Possible
+            keys are:
+                | "LOG_PROGRESS" (NewtonSolver.log_progress)
+                | "LOG_LEVEL" (NewtonSolver.log_level)
+                | "TIME_LIMIT" (NewtonSolver.time_limit)
+                | "MAXITER" (NewtonSolver.maxiter)
+                | "TOL" (NewtonSolver.tol)
+                | "BT_RHO" (NewtonSolver.rho)
+                | "BT_MAXITER" (NewtonSolver.bt_maxiter)
+                | "BACKTRACKING" (NewtonSolver.bt)
+                | "BT_START_ITER" (NewtonSolver.bt_start_iter)
+        """
         if options is None:
             options = {}
         self._options = options
 
-        if 'MAXITER' not in self._options:
+        if "LOG_PROGRESS" not in self._options:
+            self.log_progress = False
+        else:
+            self.log_progress = self._options["LOG_PROGRESS"]
+
+        if "LOG_LEVEL" not in self._options:
+            self.log_level = logging.DEBUG
+        else:
+            self.log_level = self._options["LOG_LEVEL"]
+
+        if "TIME_LIMIT" not in self._options:
+            self.time_limit = 3600
+        else:
+            self.time_limit = self._options["TIME_LIMIT"]
+
+        if "MAXITER" not in self._options:
             self.maxiter = 3000
         else:
-            self.maxiter = self._options['MAXITER']
+            self.maxiter = self._options["MAXITER"]
 
-        if 'TOL' not in self._options:
+        if "TOL" not in self._options:
             self.tol = 1e-6
         else:
-            self.tol = self._options['TOL']
+            self.tol = self._options["TOL"]
 
-        if 'BT_RHO' not in self._options:
+        if "BT_RHO" not in self._options:
             self.rho = 0.5
         else:
-            self.rho = self._options['BT_RHO']
+            self.rho = self._options["BT_RHO"]
 
-        if 'BT_MAXITER' not in self._options:
+        if "BT_MAXITER" not in self._options:
             self.bt_maxiter = 100
         else:
-            self.bt_maxiter = self._options['BT_MAXITER']
+            self.bt_maxiter = self._options["BT_MAXITER"]
 
-        if 'BACKTRACKING' not in self._options:
+        if "BACKTRACKING" not in self._options:
             self.bt = True
         else:
-            self.bt = self._options['BACKTRACKING']
+            self.bt = self._options["BACKTRACKING"]
 
-        if 'BT_START_ITER' not in self._options:
+        if "BT_START_ITER" not in self._options:
             self.bt_start_iter = 0
         else:
-            self.bt_start_iter = self._options['BT_START_ITER']
-
-        if 'THREADS' not in self._options:
-            self.num_threads = 4
-        else:
-            self.num_threads = self._options['THREADS']
+            self.bt_start_iter = self._options["BT_START_ITER"]
 
-    def solve(self, model):
+    def solve(self, model, ostream=None):
         """
 
         Parameters
@@ -71,63 +124,96 @@ def solve(self, model):
         -------
         status: SolverStatus
         message: str
+        iter_count: int
         """
-        logger_level = logger.getEffectiveLevel()
+        t0 = time.time()
 
         x = model.get_x()
         if len(x) == 0:
-            return SolverStatus.converged, 'No variables or constraints', 0
+            return (
+                SolverStatus.converged,
+                "No variables or constraints",
+                0,
+            )
 
         use_r_ = False
 
         # MAIN NEWTON LOOP
         for outer_iter in range(self.maxiter):
+            if time.time() - t0 >= self.time_limit:
+                return (
+                    SolverStatus.error,
+                    "Time limit exceeded",
+                    outer_iter,
+                )
+
             if use_r_:
                 r = r_
                 r_norm = new_norm
             else:
-                r = model.evaluate_residuals(num_threads=self.num_threads)
+                r = model.evaluate_residuals()
                 r_norm = np.max(abs(r))
 
-            if logger_level <= 1:
+            if self.log_progress or ostream is not None:
                 if outer_iter < self.bt_start_iter:
-                    logger.log(1, 'iter: {0:<4d} norm: {1:<10.2e}'.format(outer_iter, r_norm))
+                    msg = f"iter: {outer_iter:<4d} norm: {r_norm:<10.2e} time: {time.time() - t0:<8.4f}"
+                    if self.log_progress:
+                        logger.log(self.log_level, msg)
+                    if ostream is not None:
+                        ostream.write(msg + "\n")
 
             if r_norm < self.tol:
-                return SolverStatus.converged, 'Solved Successfully', outer_iter
+                return (
+                    SolverStatus.converged,
+                    "Solved Successfully",
+                    outer_iter,
+                )
 
             J = model.evaluate_jacobian(x=None)
 
             # Call Linear solver
             try:
-                d = -sp.linalg.spsolve(J, r, permc_spec='COLAMD', use_umfpack=False)
+                d = -sp.linalg.spsolve(J, r, permc_spec="COLAMD", use_umfpack=False)
             except sp.linalg.MatrixRankWarning:
-                return SolverStatus.error, 'Jacobian is singular at iteration ' + str(outer_iter), outer_iter
+                return (
+                    SolverStatus.error,
+                    "Jacobian is singular at iteration " + str(outer_iter),
+                    outer_iter,
+                )
 
             # Backtracking
             alpha = 1.0
             if self.bt and outer_iter >= self.bt_start_iter:
                 use_r_ = True
                 for iter_bt in range(self.bt_maxiter):
-                    x_ = x + alpha*d
+                    x_ = x + alpha * d
                     model.load_var_values_from_x(x_)
-                    r_ = model.evaluate_residuals(num_threads=self.num_threads)
+                    r_ = model.evaluate_residuals()
                     new_norm = np.max(abs(r_))
-                    if new_norm < (1.0-0.0001*alpha)*r_norm:
+                    if new_norm < (1.0 - 0.0001 * alpha) * r_norm:
                         x = x_
                         break
                     else:
-                        alpha = alpha*self.rho
-
-                if iter_bt+1 >= self.bt_maxiter:
-                    return SolverStatus.error, 'Line search failed at iteration ' + str(outer_iter), outer_iter
-                if logger_level <= 1:
-                    logger.log(1, 'iter: {0:<4d} norm: {1:<10.2e} alpha: {2:<10.2e}'.format(outer_iter, new_norm, alpha))
+                        alpha = alpha * self.rho
+
+                if iter_bt + 1 >= self.bt_maxiter:
+                    return (
+                        SolverStatus.error,
+                        "Line search failed at iteration " + str(outer_iter),
+                        outer_iter,
+                    )
+                if self.log_progress or ostream is not None:
+                    msg = f"iter: {outer_iter:<4d} norm: {new_norm:<10.2e} alpha: {alpha:<10.2e} time: {time.time() - t0:<8.4f}"
+                    if self.log_progress:
+                        logger.log(self.log_level, msg)
+                    if ostream is not None:
+                        ostream.write(msg + "\n")
             else:
                 x += d
                 model.load_var_values_from_x(x)
-
-        return SolverStatus.error, 'Reached maximum number of iterations: ' + str(outer_iter), outer_iter
-
-
 
+        return (
+            SolverStatus.error,
+            "Reached maximum number of iterations: " + str(outer_iter),
+            outer_iter,
+        )