Coverage for pySDC/tutorial/step_9/paradiag_setup.py: 100%

1""" 

2The ParaDiag setup of tutorial step 9, parts D and E: the advection problem of part C, a few fixed alphas and the 

3adaptive one, and a `run` that takes a communicator or not. Part D runs it with the virtually parallel controller, 

4part E on MPI ranks, which is the whole point: the same function, serial or parallel. 

5""" 

6 

7# we always do this many time-steps in total, no matter how many of them run in parallel 

8num_steps_total = 4 

9 

10 

11def get_description(): 

12 """ 

13 Set up the same advection problem as in Part C. 

14 

15 Returns: 

16 dict: the description for the ParaDiag controller 

17 """ 

18 from pySDC.implementations.problem_classes.AdvectionEquation_ND_FD import advectionNd 

19 from pySDC.implementations.sweeper_classes.ParaDiagSweepers import QDiagonalization 

20 

21 level_params = {} 

22 level_params['dt'] = 0.1 

23 level_params['restol'] = 1e-6 

24 

25 sweeper_params = {} 

26 sweeper_params['quad_type'] = 'RADAU-RIGHT' 

27 sweeper_params['num_nodes'] = 3 

28 sweeper_params['initial_guess'] = 'copy' 

29 

30 # Part C uses GMRES here to count linear solver work. We only care about the parallelism, and the 

31 # complex shifted systems ParaDiag produces are hard for GMRES, so we solve them directly instead. 

32 problem_params = {'nvars': 64, 'order': 8, 'c': 1, 'solver_type': 'direct'} 

33 

34 step_params = {} 

35 step_params['maxiter'] = 99 

36 

37 description = {} 

38 description['problem_class'] = advectionNd 

39 description['problem_params'] = problem_params 

40 description['sweeper_class'] = QDiagonalization 

41 description['sweeper_params'] = sweeper_params 

42 description['level_params'] = level_params 

43 description['step_params'] = step_params 

44 

45 return description 

46 

47 

48# the fixed values we compare against, plus the adaptive strategy 

49alpha_settings = [1e-2, 1e-4, 1e-8, 'adaptive'] 

50 

51 

52def get_controller_params(alpha): 

53 """ 

54 Controller parameters for one alpha setting. 

55 

56 Args: 

57 alpha: a number, or the string 'adaptive' 

58 

59 Returns: 

60 tuple: the controller parameters and the extra description entries 

61 """ 

62 from pySDC.implementations.convergence_controller_classes.adaptive_alpha import AdaptiveAlpha 

63 

64 controller_params = {} 

65 controller_params['logger_level'] = 30 

66 controller_params['average_jacobian'] = False 

67 

68 extra_description = {} 

69 if alpha == 'adaptive': 

70 # the adaptive controller overwrites this from the first iteration onwards, but ParaDiag needs 

71 # some alpha to build its first transform with 

72 controller_params['alpha'] = 1e-4 

73 extra_description['convergence_controllers'] = {AdaptiveAlpha: {}} 

74 else: 

75 controller_params['alpha'] = alpha 

76 

77 return controller_params, extra_description 

78 

79 

80def format_result(mode, alpha, niter, error, final_alpha): 

81 """ 

82 One line of output, in the same shape for both controllers so they can be compared. 

83 

84 Args: 

85 mode (str): which controller produced it, e.g. 'virtual' or 'MPI on 4' 

86 alpha: the alpha setting used 

87 niter (int): number of iterations needed 

88 error (float): error against the exact solution 

89 final_alpha (float): the alpha in use when the run finished 

90 

91 Returns: 

92 str: the formatted line 

93 """ 

94 return ( 

95 f'{mode:>11s}: alpha {str(alpha):>9s} -> {niter:2d} iterations, ' 

96 f'error {error:.4e}, final alpha {final_alpha:.3e}' 

97 ) 

98 

99 

100def run(alpha, block_size, comm=None): 

101 """ 

102 Run the advection problem with one alpha setting. 

103 

104 Args: 

105 alpha: a number, or the string 'adaptive' 

106 block_size (int): number of time-steps in one block 

107 comm: MPI communicator, or None for the virtually parallel controller 

108 

109 Returns: 

110 tuple: the end value, the iteration count, the error and the final alpha 

111 """ 

112 import numpy as np 

113 from pySDC.helpers.stats_helper import get_sorted 

114 

115 controller_params, extra_description = get_controller_params(alpha) 

116 description = {**get_description(), **extra_description} 

117 

118 if comm is None: 

119 from pySDC.implementations.controller_classes.controller_ParaDiag_nonMPI import controller_ParaDiag_nonMPI 

120 

121 controller_params['mssdc_jac'] = False 

122 controller = controller_ParaDiag_nonMPI( 

123 controller_params=controller_params, description=description, num_procs=block_size 

124 ) 

125 steps = controller.MS 

126 else: 

127 from pySDC.implementations.controller_classes.controller_ParaDiag_MPI import controller_ParaDiag_MPI 

128 

129 controller = controller_ParaDiag_MPI(controller_params=controller_params, description=description, comm=comm) 

130 steps = [controller.S] 

131 

132 # ParaDiag diagonalizes in time, so the solution becomes complex 

133 for S in steps: 

134 S.levels[0].prob.init = tuple([*S.levels[0].prob.init[:2]] + [np.dtype('complex128')]) 

135 

136 P = steps[0].levels[0].prob 

137 dt = steps[0].levels[0].params.dt 

138 Tend = num_steps_total * dt 

139 

140 uend, stats = controller.run(u0=P.u_exact(0.0), t0=0.0, Tend=Tend) 

141 niter = max(int(me[1]) for me in get_sorted(stats, type='niter', sortby='time', comm=comm)) 

142 

143 return uend, niter, abs(uend - P.u_exact(Tend)), controller.params.alpha