Coverage for pySDC/tutorial/step_9/C_paradiag_in_pySDC.py: 100%
95 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1# ---
2# jupyter:
3# jupytext:
4# formats: py:percent
5# kernelspec:
6# display_name: Python 3
7# name: python3
8# ---
10# %% [markdown]
11# # Part C: ParaDiag in pySDC
12#
13# Here we leave the hand-written linear algebra behind and set ParaDiag up through pySDC's controllers, comparing it
14# with single-level PFASST in Jacobi mode and with serial time stepping. In PFASST, we use a diagonal preconditioner,
15# which allows for the same amount of parallelism as ParaDiag. We show iteration counts per step, but both schemes
16# have further concurrency across the nodes.
17#
18# Two examples: a linear advection problem, discretized with finite differences, and the nonlinear van der Pol
19# oscillator, with `mu` chosen such that the problem is not overly stiff. Neither setup is optimized: with a
20# different choice of $\alpha$ in ParaDiag, or with inexactness and coarsening in PFASST, both schemes could be
21# improved significantly. This is not meant to show that one parallelization scheme is better than the other. It
22# does show that both, without optimization, need fewer iterations per task than serial time stepping. Kindly refrain
23# from computing parallel efficiency from these numbers, though. ;)
24#
25# ## The setups
26#
27# ParaDiag needs its own sweeper, `QDiagonalization`, and its own controller, `controller_ParaDiag_nonMPI`. Its
28# controller parameters set $\alpha$, and whether to average the Jacobian (for the nonlinear problem only, as it
29# costs communication).
31# %%
32import matplotlib.pyplot as plt
33import numpy as np
35from pySDC.helpers.stats_helper import get_sorted
38def get_description(problem='advection', mode='ParaDiag'):
39 level_params = {}
40 level_params['dt'] = 0.1
41 level_params['restol'] = 1e-6
43 sweeper_params = {}
44 sweeper_params['quad_type'] = 'RADAU-RIGHT'
45 sweeper_params['num_nodes'] = 3
46 sweeper_params['initial_guess'] = 'copy'
48 if mode == 'ParaDiag':
49 from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization as sweeper_class
51 # we only want to use the averaged Jacobian and do only one Newton iteration per ParaDiag iteration!
52 else:
53 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class
55 # need diagonal preconditioner for same concurrency as ParaDiag
56 sweeper_params['QI'] = 'MIN-SR-S'
58 if problem == 'advection':
59 from pySDC.implementations.problem_classes.AdvectionEquation_ND_FD import advectionNd as problem_class
61 problem_params = {'nvars': 64, 'order': 8, 'c': 1, 'solver_type': 'GMRES', 'lintol': 1e-8}
62 elif problem == 'vdp':
63 from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol as problem_class
65 # need to not raise an error when Newton has not converged because we do only one iteration
66 problem_params = {'newton_maxiter': 99, 'crash_at_maxiter': False, 'mu': 1, 'newton_tol': 1e-9}
68 step_params = {}
69 step_params['maxiter'] = 99
71 description = {}
72 description['problem_class'] = problem_class
73 description['problem_params'] = problem_params
74 description['sweeper_class'] = sweeper_class
75 description['sweeper_params'] = sweeper_params
76 description['level_params'] = level_params
77 description['step_params'] = step_params
79 return description
82def get_controller_params(problem='advection', mode='ParaDiag'):
83 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun
84 from pySDC.implementations.hooks.log_work import LogWork, LogSDCIterations
86 controller_params = {}
87 controller_params['logger_level'] = 30
88 controller_params['hook_class'] = [LogGlobalErrorPostRun, LogWork, LogSDCIterations]
90 if mode == 'ParaDiag':
91 controller_params['alpha'] = 1e-4
93 # For nonlinear problems, we need to communicate the average solution, which allows to compute the average
94 # Jacobian locally. For linear problems, we do not want the extra communication.
95 if problem == 'advection':
96 controller_params['average_jacobian'] = False
97 elif problem == 'vdp':
98 controller_params['average_jacobian'] = True
99 else:
100 # We do Block-Jacobi multi-step SDC here. It's a bit silly but it's better for comparing "speedup"
101 controller_params['mssdc_jac'] = True
103 return controller_params
106def run_problem(
107 n_steps=4,
108 problem='advection',
109 mode='ParaDiag',
110):
111 if mode == 'ParaDiag':
112 from pySDC.implementations.controller_classes.controller_ParaDiag_nonMPI import (
113 controller_ParaDiag_nonMPI as controller_class,
114 )
115 else:
116 from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI as controller_class
118 if mode == 'serial':
119 num_procs = 1
120 else:
121 num_procs = n_steps
123 description = get_description(problem, mode)
124 controller_params = get_controller_params(problem, mode)
126 controller = controller_class(num_procs=num_procs, description=description, controller_params=controller_params)
128 for S in controller.MS:
129 S.levels[0].prob.init = tuple([*S.levels[0].prob.init[:2]] + [np.dtype('complex128')])
131 P = controller.MS[0].levels[0].prob
133 t0 = 0.0
134 uinit = P.u_exact(t0)
136 uend, stats = controller.run(u0=uinit, t0=t0, Tend=n_steps * controller.MS[0].levels[0].dt)
137 return uend, stats
140# %% [markdown]
141# The solution becomes complex, because the diagonalization is: `run_problem` switches the data type of the problems
142# to complex numbers before the run.
143#
144# ## The comparison
145#
146# ParaDiag, PFASST in Jacobi mode and serial time stepping, for 16 steps. Besides the iterations, we compare the work
147# of the inner solvers: GMRES iterations for advection, Jacobian solves for van der Pol.
150# %%
151def compare_ParaDiag_and_PFASST(n_steps, problem):
152 print(f'Running {problem} with {n_steps} steps')
154 uend_PD, stats_PD = run_problem(n_steps, problem, mode='ParaDiag')
155 uend_PF, stats_PF = run_problem(n_steps, problem, mode='PFASST')
156 uend_S, stats_S = run_problem(n_steps, problem, mode='serial')
158 assert np.allclose(uend_PD, uend_PF)
159 assert np.allclose(uend_S, uend_PD)
160 assert (
161 abs(uend_PD - uend_PF) > 0
162 ) # two different iterative methods should not give identical results for non-zero tolerance
164 k_PD = get_sorted(stats_PD, type='k')
165 k_PF = get_sorted(stats_PF, type='k')
167 print(
168 f'Needed {max(me[1] for me in k_PD)} ParaDiag iterations and {max(me[1] for me in k_PF)} single-level PFASST iterations'
169 )
170 if problem == 'advection':
171 k_GMRES_PD = get_sorted(stats_PD, type='work_GMRES')
172 k_GMRES_PF = get_sorted(stats_PF, type='work_GMRES')
173 k_GMRES_S = get_sorted(stats_S, type='work_GMRES')
174 print(
175 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'
176 )
177 elif problem == 'vdp':
178 k_Jac_PD = get_sorted(stats_PD, type='work_jacobian_solves')
179 k_Jac_PF = get_sorted(stats_PF, type='work_jacobian_solves')
180 k_Jac_S = get_sorted(stats_S, type='work_jacobian_solves')
181 print(
182 f'Maximum Jacobian 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'
183 )
184 print()
185 return max(me[1] for me in k_PD), max(me[1] for me in k_PF)
188# %%
189iterations = {problem: compare_ParaDiag_and_PFASST(n_steps=16, problem=problem) for problem in ['advection', 'vdp']}
191# %% tags=["hide-input"]
192fig, ax = plt.subplots(figsize=(6, 3))
193x = np.arange(len(iterations))
194ax.bar(x - 0.2, [k[0] for k in iterations.values()], width=0.4, label='ParaDiag')
195ax.bar(x + 0.2, [k[1] for k in iterations.values()], width=0.4, label='single-level PFASST')
196ax.set_xticks(x)
197ax.set_xticklabels(['advection', 'van der Pol'])
198ax.set_ylabel('iterations')
199ax.legend(frameon=False)
200fig.tight_layout()
202# %% [markdown]
203# ParaDiag converges in very few iterations for the hyperbolic advection problem (3), where PFASST struggles (36).
204# For van der Pol, ParaDiag needs fewer iterations as well (10 against 24), with a much smaller margin. Remember
205# that ParaDiag does only one Newton iteration per ParaDiag iteration, so per node, its number of Newton iterations
206# equals the number of ParaDiag iterations, while PFASST solves the systems to some accuracy in every iteration. That
207# makes the difference in Jacobian solves per step much larger than the one in iterations: 30 against 143. Again,
208# inexactness could improve PFASST.
209#
210# :::{admonition} Important things to note
211# - ParaDiag needs its own sweeper (`QDiagonalization`) and its own controller.
212# - The solution becomes complex, because the diagonalization is.
213# - ParaDiag converges in very few iterations for the hyperbolic advection example, where PFASST struggles. For the
214# van der Pol oscillator, the gap in iterations is much smaller.
215# :::
216#
217# The checks the tests run are inside `compare_ParaDiag_and_PFASST`: all three methods agree, and the two iterative
218# ones are not identical.