Serial implementation for ParaDiag with collocation methods

pySDC/core/problem.py

@@ -164,3 +164,35 @@ def plot(self, u, t=None, fig=None):
         None
         """
         raise NotImplementedError
+
+    def solve_system(self, rhs, dt, u0, t):
+        """
+        Perform an Euler step.
+
+        Args:
+            rhs: Right hand side for the Euler step
+            dt (float): Step size for the Euler step
+            u0: Initial guess
+            t (float): Current time
+
+        Returns:
+            solution to the Euler step
+        """
+        raise NotImplementedError
+
+    def solve_jacobian(self, rhs, dt, u=None, u0=None, t=0, **kwargs):
+        """
+        Solve the Jacobian for an Euler step, linearized around u.
+        This defaults to an Euler step to accommodate linear problems.
+
+        Args:
+            rhs: Right hand side for the Euler step
+            dt (float): Step size for the Euler step
+            u: Solution to linearize around
+            u0: Initial guess
+            t (float): Current time
+
+        Returns:
+            Solution
+        """
+        return self.solve_system(rhs, dt, u0=u, t=t, **kwargs)

pySDC/implementations/controller_classes/controller_ParaDiag_nonMPI.py

@@ -213,12 +213,12 @@ def it_ParaDiag(self, local_MS_running):
             for hook in self.hooks:
                 hook.pre_sweep(step=S, level_number=0)
 
-        # replace the values stored in the steps with the residuals in order to compute the increment
-        self.swap_solution_for_all_at_once_residual(local_MS_running)
-
         # communicate average residual for setting up Jacobians for non-linear problems
         self.prepare_Jacobians(local_MS_running)
 
+        # replace the values stored in the steps with the residuals in order to compute the increment
+        self.swap_solution_for_all_at_once_residual(local_MS_running)
+
         # weighted FFT in time
         self.FFT_in_time()
 

pySDC/implementations/problem_classes/Van_der_Pol_implicit.py

@@ -69,6 +69,7 @@ def __init__(
             localVars=locals(),
         )
         self.work_counters['newton'] = WorkCounter()
+        self.work_counters['jacobian_solves'] = WorkCounter()
         self.work_counters['rhs'] = WorkCounter()
 
     def u_exact(self, t, u_init=None, t_init=None):
@@ -167,13 +168,7 @@ def solve_system(self, rhs, dt, u0, t):
             if res < self.newton_tol or np.isnan(res):
                 break
 
-            # prefactor for dg/du
-            c = 1.0 / (-2 * dt**2 * mu * x1 * x2 - dt**2 - 1 + dt * mu * (1 - x1**2))
-            # assemble dg/du
-            dg = c * np.array([[dt * mu * (1 - x1**2) - 1, -dt], [2 * dt * mu * x1 * x2 + dt, -1]])
-
-            # newton update: u1 = u0 - g/dg
-            u -= np.dot(dg, g)
+            u -= self.solve_jacobian(g, dt, u)
 
             # set new values and increase iteration count
             x1 = u[0]
@@ -191,3 +186,16 @@ def solve_system(self, rhs, dt, u0, t):
             raise ProblemError('Newton did not converge after %i iterations, error is %s' % (n, res))
 
         return u
+
+    def solve_jacobian(self, rhs, dt, u, **kwargs):
+        mu = self.mu
+        u1 = u[0]
+        u2 = u[1]
+
+        # assemble prefactor
+        c = 1.0 / (-2 * dt**2 * mu * u1 * u2 - dt**2 - 1 + dt * mu * (1 - u1**2))
+        # assemble (dg/du)^-1
+        dg = c * np.array([[dt * mu * (1 - u1**2) - 1, -dt], [2 * dt * mu * u1 * u2 + dt, -1]])
+
+        self.work_counters['jacobian_solves']()
+        return np.dot(dg, rhs)

pySDC/implementations/sweeper_classes/ParaDiagSweepers.py

@@ -105,10 +105,16 @@ def update_nodes(self):
             x1 = self.mat_vec(self.S_inv, [self.level.u[m + 1] for m in range(M)])
         else:
             x1 = self.mat_vec(self.S_inv, [self.level.u[0] for _ in range(M)])
+
+        # get averaged state over all nodes for constructing the Jacobian
+        u_avg = P.u_init
+        if not any(me is None for me in L.u_avg):
+            for m in range(M):
+                u_avg += L.u_avg[m] / M
+
         x2 = []
         for m in range(M):
-            u0 = L.u_avg[m] if L.u_avg[m] is not None else x1[m]
-            x2.append(P.solve_system(x1[m], self.w[m] * L.dt, u0=u0, t=L.time + L.dt * self.coll.nodes[m]))
+            x2.append(P.solve_jacobian(x1[m], self.w[m] * L.dt, u=u_avg, t=L.time + L.dt * self.coll.nodes[m]))
         z = self.mat_vec(self.S, x2)
         y = self.mat_vec(self.params.G_inv, z)
 

pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py

@@ -15,11 +15,10 @@
 have the same Jacobian on all steps.
 The ParaDiag iteration then proceeds as follows:
     - (1) Compute residual of composite collocation problem
-    - (2) Average the residual across the steps as preparation for computing the average Jacobian
-          Note that we still have values for each collocation node and space position.
+    - (2) Average the solution across the steps and nodes as preparation for computing the average Jacobian
     - (3) Weighted FFT in time to diagonalize E_alpha
-    - (4) Solve for the increment by perform a single Newton iterations on the subproblems on the different steps and
-          nodes. The Jacobian is based on the averaged residual from (2)
+    - (4) Solve for the increment by inverting the averaged Jacobian from (2) on the subproblems on the different steps
+          and nodes.
     - (5) Weighted iFFT in time
     - (6) Increment solution
 As IMEX ParaDiag is a trivial extension of ParaDiag for linear problems, we focus on the second approach here.
@@ -170,15 +169,15 @@ def residual(_u, u0):
 sol_paradiag = u.copy() * 0j
 u0 = u.copy()
 
-buf = prob.u_init
 niter = 0
 res = residual(sol_paradiag, u0)
 while np.max(np.abs(res)) > restol:
     # compute all-at-once residual
     res = residual(sol_paradiag, u0)
 
-    # compute residual averaged across the L steps. This is the difference to ParaDiag for linear problems.
-    res_avg = np.mean(res, axis=0)
+    # compute residual averaged across the L steps and M nodes. This is the difference to ParaDiag for linear problems.
+    u_avg = prob.u_init
+    u_avg[:] = np.mean(sol_paradiag, axis=(0, 1))
 
     # weighted FFT in time
     x = np.fft.fft(mat_vec(J_inv.toarray(), res), axis=0)
@@ -191,8 +190,7 @@ def residual(_u, u0):
         x1 = S_inv[l] @ x[l]
         x2 = np.empty_like(x1)
         for m in range(M):
-            buf[:] = res_avg[m]  # set up averaged Jacobian by using averaged residual as initial guess
-            x2[m, :] = prob.solve_system(x1[m], w[l][m] * dt, u0=buf, t=l * dt)
+            x2[m, :] = prob.solve_jacobian(x1[m], w[l][m] * dt, u=u_avg, t=l * dt)
         z = S[l] @ x2
         y[l, ...] = sp.linalg.spsolve(G[l], z)
 

pySDC/tutorial/step_9/C_paradiag_in_pySDC.py

@@ -47,11 +47,9 @@ def get_description(problem='advection', mode='ParaDiag'):
         from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization as sweeper_class
 
         # we only want to use the averaged Jacobian and do only one Newton iteration per ParaDiag iteration!
-        newton_maxiter = 1
     else:
         from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
 
-        newton_maxiter = 99
         # need diagonal preconditioner for same concurrency as ParaDiag
         sweeper_params['QI'] = 'MIN-SR-S'
 
@@ -63,7 +61,7 @@ def get_description(problem='advection', mode='ParaDiag'):
         from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol as problem_class
 
         # need to not raise an error when Newton has not converged because we do only one iteration
-        problem_params = {'newton_maxiter': newton_maxiter, 'crash_at_maxiter': False, 'mu': 1, 'newton_tol': 1e-9}
+        problem_params = {'newton_maxiter': 99, 'crash_at_maxiter': False, 'mu': 1, 'newton_tol': 1e-9}
 
     step_params = {}
     step_params['maxiter'] = 99
@@ -164,11 +162,11 @@ def compare_ParaDiag_and_PFASST(n_steps, problem):
             f'Maximum GMRES iterations on each step: {max(me[1] for me in k_GMRES_PD)} in ParaDiag, {max(me[1] for me in k_GMRES_PF)} in single-level PFASST and {sum(me[1] for me in k_GMRES_S)} total GMRES iterations in serial'
         )
     elif problem == 'vdp':
-        k_Newton_PD = get_sorted(stats_PD, type='work_newton')
-        k_Newton_PF = get_sorted(stats_PF, type='work_newton')
-        k_Newton_S = get_sorted(stats_S, type='work_newton')
+        k_Jac_PD = get_sorted(stats_PD, type='work_jacobian_solves')
+        k_Jac_PF = get_sorted(stats_PF, type='work_jacobian_solves')
+        k_Jac_S = get_sorted(stats_S, type='work_jacobian_solves')
         my_print(
-            f'Maximum Newton iterations on each step: {max(me[1] for me in k_Newton_PD)} in ParaDiag, {max(me[1] for me in k_Newton_PF)} in single-level PFASST and {sum(me[1] for me in k_Newton_S)} total Newton iterations in serial'
+            f'Maximum Jacabian solves on each step: {max(me[1] for me in k_Jac_PD)} in ParaDiag, {max(me[1] for me in k_Jac_PF)} in single-level PFASST and {sum(me[1] for me in k_Jac_S)} total Jacobian solves in serial'
         )
     my_print()