Added a test for counting Newton iterations

Author: Thomas Baumann

pySDC/implementations/hooks/log_work.py

@@ -1,5 +1,6 @@
 from pySDC.core.Hooks import hooks
 
+
 class LogWork(hooks):
     """
     Log the increment of all work counters in the problem between steps
@@ -18,7 +19,8 @@ def pre_step(self, step, level_number):
         """
         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))
+                {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):

pySDC/implementations/problem_classes/Van_der_Pol_implicit.py

@@ -39,7 +39,6 @@ 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.total_newton_iter = 0  # access this from the hooks to make stats
 
     def u_exact(self, t, u_init=None, t_init=None):
         """
@@ -129,7 +128,6 @@ def solve_system(self, rhs, dt, u0, t):
             x1 = u[0]
             x2 = u[1]
             n += 1
-        self.total_newton_iter += n
 
         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,56 +8,6 @@
 hook_collection = [LogSolution, LogEmbeddedErrorEstimate, LogExtrapolationErrorEstimate, LogStepSize]
 
 
-class LogSpaceIter(hooks):
-    """
-    Log the number of iterations from the space solver e.g. CG or GMRES after each step
-    """
-
-    def pre_step(self, step, level_number):
-        if level_number == 0:
-            if 'space_iter_counter' in step.levels[0].prob.__dict__.keys():
-                self.__space_iter_last_step = [
-                    step.levels[i].prob.space_iter_counter.niter for i in range(len(step.levels))
-                ]
-
-    def post_step(self, step, level_number):
-        L = step.levels[level_number]
-        if 'space_iter_counter' in L.prob.__dict__.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='space_iter',
-                value=L.prob.space_iter_counter.niter - self.__space_iter_last_step[level_number],
-            )
-
-
-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,23 @@ 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 = {}
 
     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 +227,16 @@ 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')]
+
+        if imex:
+            assert all([me == 0 for me in newton_iter]), f"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)
 
@@ -236,100 +251,6 @@ def compare_imex_full(plotting=False):
     ), f"Expected runaway to happen, but maximum temperature is {max(res[True]):.2e} < u_max={prob.params.u_max:.2e}!"
 
 
-def compare_dirk():
-    from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity, AdaptivityRK
-    from pySDC.implementations.sweeper_classes.Runge_Kutta import DIRK34
-    from pySDC.projects.Resilience.hook import hook_collection, LogNewtonIter, LogData, LogSpaceIter
-    from pySDC.helpers.plot_helper import figsize_by_journal, setup_mpl
-    from pySDC.implementations.hooks.log_work import LogWork
-
-    setup_mpl(reset=True)
-    fig, axs = plt.subplots(2, 2, figsize=figsize_by_journal('Springer_Numerical_Algorithms', 1, 0.82), sharex=True)
-
-    problem_params = {}
-    problem_params['direct_solver'] = False
-
-    description = {}
-    description['step_params'] = {'maxiter': 4}
-    description['convergence_controllers'] = {Adaptivity: {'e_tol': 1e-5, 'dt_max': 50.0}}
-
-    description_RK = {}
-    description_RK['convergence_controllers'] = {AdaptivityRK: {'e_tol': 1e-5, 'dt_max': 5e1, 'update_order': 4}}
-    description_RK['sweeper_class'] = DIRK34
-    description_RK['step_params'] = {'maxiter': 1}
-
-    description_res = {}
-    description_res['step_params'] = {'maxiter': 99}
-    description_res['level_params'] = {'restol': 1e-10}
-
-    controller_params = {}
-    controller_params['hook_class'] = hook_collection + [LogNewtonIter, LogData, LogSpaceIter, LogWork]
-    controller_params['logger_level'] = 30
-    Tend = 500
-
-    stats_fixed, _, _ = run_leaky_superconductor(
-        custom_description=description_res,
-        imex=False,
-        Tend=Tend,
-        custom_controller_params=controller_params,
-        custom_problem_params=problem_params,
-    )
-    stats, _, _ = run_leaky_superconductor(
-        custom_description=description,
-        imex=False,
-        Tend=Tend,
-        custom_controller_params=controller_params,
-        custom_problem_params=problem_params,
-    )
-    stats_RK, _, _ = run_leaky_superconductor(
-        custom_description=description_RK,
-        imex=False,
-        Tend=Tend,
-        custom_controller_params=controller_params,
-        custom_problem_params=problem_params,
-    )
-
-    labels = {
-        0: 'SDC+adaptivity',
-        1: 'DIRK4(3)',
-        2: 'SDC',
-    }
-    ls = {
-        0: '-',
-        1: '--',
-        2: ':',
-    }
-    S = [stats, stats_RK, stats_fixed]
-    for i in range(len(S)):
-        newton_iter = get_sorted(S[i], type='work_newton')
-        axs[0, 0].plot(
-            [me[0] for me in newton_iter], np.cumsum([me[1] for me in newton_iter]), label=labels[i], ls=ls[i]
-        )
-
-        space_iter = get_sorted(S[i], type='work_linear')
-        axs[1, 1].plot([me[0] for me in space_iter], np.cumsum([me[1] for me in space_iter]), ls=ls[i])
-
-        dt = get_sorted(S[i], type='dt')
-        axs[0, 1].plot([me[0] for me in dt], [me[1] for me in dt], ls=ls[i])
-
-        # u = get_sorted(S[i], type='u', recomputed=False)
-        # axs[1,1].plot([me[0] for me in u], [max(me[1]) for me in u], ls=ls[i])
-
-        k = get_sorted(S[i], type='sweeps', recomputed=None)
-        axs[1, 0].plot([me[0] for me in k], np.cumsum([me[1] for me in k]), ls=ls[i])
-    axs[0, 0].set_yscale('log')
-    axs[1, 1].set_yscale('log')
-    axs[0, 0].legend(frameon=False)
-    axs[0, 0].set_ylabel('cumulative Newton iterations')
-    axs[1, 0].set_ylabel('cumulative SDC iterations')
-    axs[1, 0].set_xlabel('$t$')
-    axs[0, 1].set_ylabel(r'$\Delta t$')
-    axs[1, 1].set_ylabel('cumulative GMRES iterations')
-    fig.tight_layout()
-    fig.savefig('data/quench_newton.pdf')
-
-
 if __name__ == '__main__':
-    compare_dirk()
-    # compare_imex_full(plotting=True)
+    compare_imex_full(plotting=True)
     plt.show()