Coverage for pySDC/tutorial/step_5/C_advection_and_PFASST.py: 100%

1# --- 

2# jupyter: 

3# jupytext: 

4# formats: py:percent 

5# kernelspec: 

6# display_name: Python 3 

7# name: python3 

8# --- 

9 

10# %% [markdown] 

11# # Part C: Advection and PFASST 

12# 

13# PFASST did very well for the heat equation in [Part B](B_my_first_PFASST_run). Now the same test for an 

14# advection problem, $u_t + c\, u_x = 0$ with periodic boundaries. The setup is the same, but with the fully 

15# implicit sweeper `generic_implicit`, and with two preconditioners to compare: the LU-trick and implicit Euler. 

16# 

17# The setup is periodic in time as well: at `Tend = 1`, the wave has travelled once through the domain, and the 

18# exact solution looks exactly like the initial condition. 

19 

20# %% 

21import matplotlib.pyplot as plt 

22import numpy as np 

23 

24from pySDC.helpers.stats_helper import get_sorted 

25from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

26from pySDC.implementations.problem_classes.AdvectionEquation_ND_FD import advectionNd 

27from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit 

28from pySDC.implementations.transfer_classes.TransferMesh import mesh_to_mesh 

29 

30# initialize level parameters 

31level_params = {'restol': 1e-09, 'dt': 0.0625} 

32 

33# initialize sweeper parameters, the preconditioner QI is set in the loop below 

34sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': [3]} 

35 

36# initialize problem parameters 

37problem_params = { 

38 'c': 1, # advection coefficient 

39 'freq': 4, # frequency for the test value 

40 'nvars': [128, 64], # number of degrees of freedom for each level 

41 'order': 4, 

42 'bc': 'periodic', 

43 'stencil_type': 'center', 

44} 

45 

46# initialize step parameters 

47step_params = {'maxiter': 50} 

48 

49# initialize space transfer parameters 

50space_transfer_params = {'rorder': 2, 'iorder': 6, 'periodic': True} 

51 

52# initialize controller parameters 

53controller_params = {'logger_level': 30, 'predict_type': 'pfasst_burnin'} 

54 

55# fill description dictionary for easy step instantiation 

56description = { 

57 'problem_class': advectionNd, # pass problem class 

58 'problem_params': problem_params, # pass problem parameters 

59 'sweeper_class': generic_implicit, # pass sweeper 

60 'level_params': level_params, # pass level parameters 

61 'step_params': step_params, # pass step parameters 

62 'space_transfer_class': mesh_to_mesh, # pass spatial transfer class 

63 'space_transfer_params': space_transfer_params, # pass parameters for spatial transfer 

64} 

65 

66# set time parameters 

67t0 = 0.0 

68Tend = 1.0 

69 

70# set up list of parallel time-steps to run PFASST with 

71nsteps = int(Tend / level_params['dt']) 

72num_proc_list = [2**i for i in range(int(np.log2(nsteps) + 1))] 

73 

74# set up list of types of implicit SDC sweepers: LU and implicit Euler here 

75QI_list = ['LU', 'IE'] 

76 

77# %% [markdown] 

78# ## The runs 

79# 

80# As in Part B, for each preconditioner and each number of parallel steps: 

81 

82# %% tags=["scroll-output"] 

83results = {} 

84# loop over different types of implicit sweeper types 

85for QI in QI_list: 

86 # define and set preconditioner for the implicit sweeper 

87 sweeper_params['QI'] = QI 

88 description['sweeper_params'] = sweeper_params # pass sweeper parameters 

89 

90 # loop over different number of processes 

91 for num_proc in num_proc_list: 

92 print('Working with QI = %s on %2i processes...' % (QI, num_proc)) 

93 # instantiate controller 

94 controller = controller_nonMPI(num_procs=num_proc, controller_params=controller_params, description=description) 

95 

96 # get initial values on finest level 

97 P = controller.MS[0].levels[0].prob 

98 uinit = P.u_exact(t0) 

99 

100 # call main function to get things done... 

101 uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend) 

102 

103 # compute exact solution and compare 

104 uex = P.u_exact(Tend) 

105 err = abs(uex - uend) 

106 

107 # filter statistics by type (number of iterations) 

108 iter_counts = get_sorted(stats, type='niter', sortby='time') 

109 

110 # compute and print statistics 

111 niters = np.array([item[1] for item in iter_counts]) 

112 print(' Mean number of iterations: %4.2f' % np.mean(niters)) 

113 print(' Range of values for number of iterations: %2i ' % np.ptp(niters)) 

114 print( 

115 ' Position of max/min number of iterations: %2i -- %2i' % (int(np.argmax(niters)), int(np.argmin(niters))) 

116 ) 

117 print( 

118 ' Std and var for number of iterations: %4.2f -- %4.2f' % (float(np.std(niters)), float(np.var(niters))) 

119 ) 

120 print() 

121 

122 results[QI, num_proc] = (err, niters) 

123 

124 means = [np.mean(results[QI, n][1]) for n in num_proc_list] 

125 print('Mean number of iterations went up from %4.2f to %4.2f for QI = %s!' % (min(means), max(means), QI)) 

126 print() 

127 print() 

128 

129# %% tags=["hide-input"] 

130fig, axes = plt.subplots(1, len(QI_list), figsize=(10, 2.6), sharey=True, constrained_layout=True) 

131vmax = max(max(niters) for _, niters in results.values()) 

132for ax, QI in zip(axes, QI_list, strict=True): 

133 image = ax.imshow([results[QI, n][1] for n in num_proc_list], cmap='viridis', aspect='auto', vmin=0, vmax=vmax) 

134 ax.set_title(f'QI = {QI}', fontsize=10) 

135 ax.set_xlabel('time step') 

136axes[0].set_yticks(range(len(num_proc_list))) 

137axes[0].set_yticklabels(num_proc_list) 

138axes[0].set_ylabel('parallel steps') 

139fig.colorbar(image, ax=axes, label='iterations') 

140 

141# %% [markdown] 

142# Unlike for the heat equation, the iteration counts grow significantly with the number of parallel steps, and 

143# within each block of parallel steps the later ones need more. How much depends on the problem, but this is typical 

144# for parallel-in-time methods of this kind. Note also that the LU-trick, so good for the heat equation, is no 

145# better than implicit Euler here, and even slightly worse with 8 and 16 parallel steps. 

146# 

147# :::{admonition} Try it yourself 

148# :class: tip 

149# Add `'MIN-SR-S'` to `QI_list`: a diagonal preconditioner, whose nodes could be solved in parallel. How does it 

150# fare with 16 parallel steps? 

151# ::: 

152# 

153# :::{dropdown} Answer 

154# Much worse: about 24 iterations per step on average, against 15 for `'LU'` and 13 for `'IE'`. With one step at a 

155# time, all three need 5. With 16 parallel steps, the last four hit `maxiter = 50` without reaching `restol`, so its 

156# error, $5.2 \cdot 10^{-4}$, is above what the check in the last cell allows. All converged runs end at the same 

157# error, set by the spatial discretization, just below that threshold. 

158# ::: 

159# 

160# :::{admonition} Important things to note 

161# - Like the IMEX sweeper, `generic_implicit` lets you choose the preconditioner with `QI`: `'IE'` for implicit 

162# Euler, `'LU'` for the LU-trick, and more. They come from the [qmat](https://github.com/Parallel-in-Time/qmat) 

163# package, and `qmat.qdelta.QDELTA_GENERATORS` lists them all, implicit and explicit ones. 

164# ::: 

165# 

166# The check the tests run, for every run: 

167 

168# %% 

169for err, _ in results.values(): 

170 assert err < 5.1365e-04, f"ERROR: error is too high, got {err}"