Added some more work counters

Author: Thomas Baumann

pySDC/implementations/problem_classes/LeakySuperconductor.py

@@ -96,6 +96,7 @@ 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.work_counters['newton'] = WorkCounter()
+        self.work_counters['rhs'] = WorkCounter()
         if not self.params.direct_solver:
             self.work_counters['linear'] = WorkCounter()
 
@@ -141,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):
@@ -196,6 +198,7 @@ def get_non_linear_Jacobian(u):
         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
@@ -281,6 +284,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/leaky_superconductor.py

@@ -209,6 +209,7 @@ def compare_imex_full(plotting=False):
     newton_iter_max = 20
 
     res = {}
+    rhs = {}
 
     custom_description = {}
     custom_description['problem_params'] = {
@@ -232,6 +233,7 @@ def compare_imex_full(plotting=False):
 
         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!"
@@ -252,6 +254,8 @@ 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)