Merge pull request #279 from brownbaerchen/work_models

Counters for measuring work

Author: pancetta

pySDC/core/Problem.py

@@ -12,6 +12,18 @@ def __init__(self, pars):
         self._freeze()
 
 
+class WorkCounter(object):
+    """
+    Class for counting iterations
+    """
+
+    def __init__(self):
+        self.niter = 0
+
+    def __call__(self, *args, **kwargs):
+        self.niter += 1
+
+
 class ptype(object):
     """
     Prototype class for problems, just defines the attributes essential to get started
@@ -36,6 +48,7 @@ def __init__(self, init, dtype_u, dtype_f, params):
         """
 
         self.params = _Pars(params)
+        self.work_counters = {}
 
         # set up logger
         self.logger = logging.getLogger('problem')

pySDC/implementations/hooks/log_work.py

@@ -0,0 +1,47 @@
+from pySDC.core.Hooks import hooks
+
+
+class LogWork(hooks):
+    """
+    Log the increment of all work counters in the problem between steps
+    """
+
+    def pre_step(self, step, level_number):
+        """
+        Store the current values of the work counters
+
+        Args:
+            step (pySDC.Step.step): the current step
+            level_number (int): the current level number
+
+        Returns:
+            None
+        """
+        if level_number == 0:
+            self.__work_last_step = [
+                {key: step.levels[i].prob.work_counters[key].niter for key in step.levels[i].prob.work_counters.keys()}
+                for i in range(len(step.levels))
+            ]
+
+    def post_step(self, step, level_number):
+        """
+        Add the difference between current values of counters and their values before the iteration to the stats.
+
+        Args:
+            step (pySDC.Step.step): the current step
+            level_number (int): the current level number
+
+        Returns:
+            None
+        """
+        L = step.levels[level_number]
+        for key in self.__work_last_step[level_number].keys():
+            self.add_to_stats(
+                process=step.status.slot,
+                time=L.time + L.dt,
+                level=L.level_index,
+                iter=step.status.iter,
+                sweep=L.status.sweep,
+                type=f'work_{key}',
+                value=L.prob.work_counters[key].niter - self.__work_last_step[level_number][key],
+            )

pySDC/implementations/problem_classes/LeakySuperconductor.py

@@ -1,9 +1,9 @@
 import numpy as np
 import scipy.sparse as sp
-from scipy.sparse.linalg import spsolve
+from scipy.sparse.linalg import spsolve, gmres, inv
 
 from pySDC.core.Errors import ParameterError, ProblemError
-from pySDC.core.Problem import ptype
+from pySDC.core.Problem import ptype, WorkCounter
 from pySDC.helpers import problem_helper
 from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 
@@ -46,6 +46,9 @@ def __init__(self, problem_params, dtype_u=mesh, dtype_f=mesh):
             'nvars': 2**7,
             'newton_tol': 1e-8,
             'newton_iter': 99,
+            'lintol': 1e-8,
+            'liniter': 99,
+            'direct_solver': True,
         }
 
         for key in problem_params.keys():
@@ -92,7 +95,10 @@ def __init__(self, problem_params, dtype_u=mesh, dtype_f=mesh):
 
         self.leak = np.logical_and(self.xv > self.params.leak_range[0], self.xv < self.params.leak_range[1])
 
-        self.total_newton_iter = 0  # store here how many iterations you needed for the Newton solver over the entire run to extract the desired information in the hooks
+        self.work_counters['newton'] = WorkCounter()
+        self.work_counters['rhs'] = WorkCounter()
+        if not self.params.direct_solver:
+            self.work_counters['linear'] = WorkCounter()
 
     def eval_f_non_linear(self, u, t):
         """
@@ -136,6 +142,7 @@ def eval_f(self, u, t):
         """
         f = self.dtype_f(self.init)
         f[:] = self.A.dot(u.flatten()).reshape(self.params.nvars) + self.eval_f_non_linear(u, t)
+        self.work_counters['rhs']()
         return f
 
     def solve_system(self, rhs, factor, u0, t):
@@ -182,9 +189,20 @@ def get_non_linear_Jacobian(u):
 
         u = self.dtype_u(u0)
         res = np.inf
+        delta = np.zeros_like(u)
+
+        # construct a preconditioner for the space solver
+        if not self.params.direct_solver:
+            M = inv(self.Id - factor * self.A)
+
         for n in range(0, self.params.newton_iter):
             # assemble G such that G(u) = 0 at the solution of the step
             G = u - factor * self.eval_f(u, t) - rhs
+            self.work_counters[
+                'rhs'
+            ].niter -= (
+                1  # Work regarding construction of the Jacobian etc. should count into the Newton iterations only
+            )
 
             res = np.linalg.norm(G, np.inf)
             if res <= self.params.newton_tol and n > 0:  # we want to make at least one Newton iteration
@@ -194,15 +212,27 @@ def get_non_linear_Jacobian(u):
             J = self.Id - factor * (self.A + get_non_linear_Jacobian(u))
 
             # solve the linear system
-            delta = spsolve(J, G)
+            if self.params.direct_solver:
+                delta = spsolve(J, G)
+            else:
+                delta, info = gmres(
+                    J,
+                    G,
+                    x0=delta,
+                    M=M,
+                    tol=self.params.lintol,
+                    maxiter=self.params.liniter,
+                    atol=0,
+                    callback=self.work_counters['linear'],
+                )
 
             if not np.isfinite(delta).all():
                 break
 
             # update solution
             u = u - delta
 
-        self.total_newton_iter += n
+            self.work_counters['newton']()
 
         return u
 
@@ -258,6 +288,7 @@ def eval_f(self, u, t):
         f.impl[:] = self.A.dot(u.flatten()).reshape(self.params.nvars)
         f.expl[:] = self.eval_f_non_linear(u, t)
 
+        self.work_counters['rhs']()
         return f
 
     def solve_system(self, rhs, factor, u0, t):

pySDC/implementations/problem_classes/Lorenz.py

@@ -1,5 +1,5 @@
 import numpy as np
-from pySDC.core.Problem import ptype
+from pySDC.core.Problem import ptype, WorkCounter
 from pySDC.implementations.datatype_classes.mesh import mesh
 
 
@@ -44,6 +44,7 @@ def __init__(self, problem_params):
             dtype_f=mesh,
             params=problem_params,
         )
+        self.work_counters['newton'] = WorkCounter()
 
     def eval_f(self, u, t):
         """
@@ -122,6 +123,7 @@ def solve_system(self, rhs, dt, u0, t):
 
             # update solution
             u = u - delta
+            self.work_counters['newton']()
 
         return u
 

pySDC/implementations/problem_classes/Van_der_Pol_implicit.py

@@ -2,7 +2,7 @@
 from scipy.integrate import solve_ivp
 
 from pySDC.core.Errors import ParameterError, ProblemError
-from pySDC.core.Problem import ptype
+from pySDC.core.Problem import ptype, WorkCounter
 from pySDC.implementations.datatype_classes.mesh import mesh
 
 
@@ -39,6 +39,7 @@ def __init__(self, problem_params, dtype_u=mesh, dtype_f=mesh):
         super(vanderpol, self).__init__(
             (problem_params['nvars'], None, np.dtype('float64')), dtype_u, dtype_f, problem_params
         )
+        self.work_counters['newton'] = WorkCounter()
 
     def u_exact(self, t, u_init=None, t_init=None):
         """
@@ -128,6 +129,7 @@ def solve_system(self, rhs, dt, u0, t):
             x1 = u[0]
             x2 = u[1]
             n += 1
+            self.work_counters['newton']()
 
         if np.isnan(res) and self.params.stop_at_nan:
             raise ProblemError('Newton got nan after %i iterations, aborting...' % n)

pySDC/projects/Resilience/fault_stats.py

@@ -9,12 +9,13 @@
 import pySDC.helpers.plot_helper as plot_helper
 from pySDC.helpers.stats_helper import get_sorted
 
-from pySDC.projects.Resilience.hook import hook_collection, LogUAllIter, LogData, LogNewtonIter
+from pySDC.projects.Resilience.hook import hook_collection, LogUAllIter, LogData
 from pySDC.projects.Resilience.fault_injection import get_fault_injector_hook
 from pySDC.implementations.convergence_controller_classes.hotrod import HotRod
 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
 from pySDC.implementations.convergence_controller_classes.basic_restarting import BasicRestartingNonMPI
 from pySDC.implementations.hooks.log_errors import LogLocalErrorPostStep
+from pySDC.implementations.hooks.log_work import LogWork
 
 # these problems are available for testing
 from pySDC.projects.Resilience.advection import run_advection
@@ -620,7 +621,7 @@ def generate_stats(self, strategy=None, runs=1000, reload=True, faults=True, com
             else:
                 error = self.get_error(u, t, controller, strategy)
             total_iteration = sum([k[1] for k in get_sorted(stats, type='k')])
-            total_newton_iteration = sum([k[1] for k in get_sorted(stats, type='newton_iter')])
+            total_newton_iteration = sum([k[1] for k in get_sorted(stats, type='work_newton')])
 
             # record the new data point
             if faults:
@@ -694,7 +695,7 @@ def single_run(self, strategy, run=0, faults=False, force_params=None, hook_clas
             pySDC.Controller: The controller of the run
             float: The time the problem should have run to
         '''
-        hook_class = hook_collection + [LogNewtonIter] + ([LogData] if hook_class is None else hook_class)
+        hook_class = hook_collection + [LogWork] + ([LogData] if hook_class is None else hook_class)
         force_params = {} if force_params is None else force_params
 
         # build the custom description
@@ -870,7 +871,7 @@ def scrutinize(self, strategy, run, faults=True):
         )
         print(f'k: sum: {np.sum(k)}, min: {np.min(k)}, max: {np.max(k)}, mean: {np.mean(k):.2f},')
 
-        _newton_iter = get_sorted(stats, type='newton_iter')
+        _newton_iter = get_sorted(stats, type='work_newton')
         if len(_newton_iter) > 0:
             newton_iter = [me[1] for me in _newton_iter]
             print(

pySDC/projects/Resilience/hook.py

@@ -8,30 +8,6 @@
 hook_collection = [LogSolution, LogEmbeddedErrorEstimate, LogExtrapolationErrorEstimate, LogStepSize]
 
 
-class LogNewtonIter(hooks):
-    """
-    Log the number of Newton iterations required for each step
-    """
-
-    def pre_step(self, step, level_number):
-        if level_number == 0:
-            self.__newton_iter_last_step = [
-                step.levels[i].prob.__dict__.get('total_newton_iter', 0) for i in range(len(step.levels))
-            ]
-
-    def post_step(self, step, level_number):
-        L = step.levels[level_number]
-        self.add_to_stats(
-            process=step.status.slot,
-            time=L.time + L.dt,
-            level=L.level_index,
-            iter=step.status.iter,
-            sweep=L.status.sweep,
-            type='newton_iter',
-            value=L.prob.__dict__.get('total_newton_iter', 0) - self.__newton_iter_last_step[level_number],
-        )
-
-
 class LogData(hooks):
     """
     Record data required for analysis of problems in the resilience project

pySDC/projects/Resilience/leaky_superconductor.py

@@ -202,15 +202,24 @@ def compare_imex_full(plotting=False):
         plotting (bool): Plot the solution or not
     """
     from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
+    from pySDC.implementations.hooks.log_work import LogWork
+
+    maxiter = 5
+    num_nodes = 3
+    newton_iter_max = 20
 
     res = {}
+    rhs = {}
 
     custom_description = {}
     custom_description['problem_params'] = {
         'newton_tol': 1e-10,
-        'newton_iter': 20,
+        'newton_iter': newton_iter_max,
         'nvars': 2**9,
     }
+    custom_description['step_params'] = {'maxiter': maxiter}
+    custom_description['sweeper_params'] = {'num_nodes': num_nodes}
+
     custom_controller_params = {'logger_level': 30}
     for imex in [False, True]:
         custom_description['convergence_controllers'] = {Adaptivity: {'e_tol': 1e-6, 'dt_max': 1e2}}
@@ -219,9 +228,19 @@ def compare_imex_full(plotting=False):
             custom_controller_params=custom_controller_params,
             imex=imex,
             Tend=5e2,
+            hook_class=LogWork,
         )
 
         res[imex] = get_sorted(stats, type='u')[-1][1]
+        newton_iter = [me[1] for me in get_sorted(stats, type='work_newton')]
+        rhs[imex] = np.mean([me[1] for me in get_sorted(stats, type='work_rhs')]) // 1
+
+        if imex:
+            assert all([me == 0 for me in newton_iter]), "IMEX is not supposed to do Newton iterations!"
+        else:
+            assert (
+                max(newton_iter) / num_nodes / maxiter <= newton_iter_max
+            ), "Took more Newton iterations than allowed!"
         if plotting:  # pragma no cover
             plot_solution(stats, controller)
 
@@ -235,6 +254,10 @@ def compare_imex_full(plotting=False):
         max(res[True]) > prob.params.u_max
     ), f"Expected runaway to happen, but maximum temperature is {max(res[True]):.2e} < u_max={prob.params.u_max:.2e}!"
 
+    assert (
+        rhs[True] == rhs[False]
+    ), f"Expected IMEX and fully implicit schemes to take the same number of right hand side evaluations per step, but got {rhs[True]} and {rhs[False]}!"
+
 
 if __name__ == '__main__':
     compare_imex_full(plotting=True)