Coverage for pySDC/tutorial/step_7/B_pySDC_with_mpi4pyfft.py: 100%

1# --- 

2# jupyter: 

3# jupytext: 

4# formats: py:percent 

5# kernelspec: 

6# display_name: Python 3 

7# name: python3 

8# language_info: 

9# name: python 

10# --- 

11 

12# %% [markdown] 

13# # Part B: mpi4py-fft for parallel Fourier transforms 

14# 

15# The most prominent parallel solver is, probably, the FFT. While many implementations or wrappers for Python exist, 

16# we decided to use [mpi4py-fft](https://mpi4py-fft.readthedocs.io/en/latest/), which provided the easiest 

17# installation, a simple API and good parallel scaling. As an example, we solve the nonlinear Schrödinger equation in 

18# 2D, with the IMEX sweeper treating the nonlinear parts explicitly. The code works both in real and in spectral 

19# space, the latter usually being faster. This example tests SDC, MLSDC and PFASST. 

20# 

21# ## One run 

22# 

23# The problem class `nonlinearschroedinger_imex` gets the MPI communicator, over which mpi4py-fft distributes the 

24# grid, and `fft_to_fft` transfers between the levels. Each run appends its error, statistics of the iterations and 

25# the time to solution to `data/step_7_B_out.txt`, from rank 0. 

26 

27# %% 

28import numpy as np 

29from pathlib import Path 

30from mpi4py import MPI 

31 

32from pySDC.helpers.stats_helper import get_sorted 

33 

34from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

35from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order 

36from pySDC.implementations.problem_classes.NonlinearSchroedinger_MPIFFT import nonlinearschroedinger_imex 

37from pySDC.implementations.transfer_classes.TransferMesh_MPIFFT import fft_to_fft 

38 

39 

40def run_simulation(spectral=None, ml=None, num_procs=None): 

41 """ 

42 A test program to do SDC, MLSDC and PFASST runs for the 2D NLS equation 

43 

44 Args: 

45 spectral (bool): run in real or spectral space 

46 ml (bool): single or multiple levels 

47 num_procs (int): number of parallel processors 

48 """ 

49 

50 comm = MPI.COMM_WORLD 

51 rank = comm.Get_rank() 

52 

53 # initialize level parameters 

54 level_params = dict() 

55 level_params['restol'] = 1e-08 

56 level_params['dt'] = 1e-01 / 2 

57 level_params['nsweeps'] = [1] 

58 

59 # initialize sweeper parameters 

60 sweeper_params = dict() 

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

62 sweeper_params['num_nodes'] = [3] 

63 sweeper_params['QI'] = ['LU'] # For the IMEX sweeper, the LU-trick can be activated for the implicit part 

64 sweeper_params['initial_guess'] = 'zero' 

65 

66 # initialize problem parameters 

67 problem_params = dict() 

68 if ml: 

69 problem_params['nvars'] = [(128, 128), (32, 32)] 

70 else: 

71 problem_params['nvars'] = [(128, 128)] 

72 problem_params['spectral'] = spectral 

73 problem_params['comm'] = comm 

74 

75 # initialize step parameters 

76 step_params = dict() 

77 step_params['maxiter'] = 50 

78 

79 # initialize controller parameters 

80 controller_params = dict() 

81 controller_params['logger_level'] = 30 if rank == 0 else 99 

82 # controller_params['predict_type'] = 'fine_only' 

83 

84 # fill description dictionary for easy step instantiation 

85 description = dict() 

86 description['problem_params'] = problem_params # pass problem parameters 

87 description['problem_class'] = nonlinearschroedinger_imex 

88 description['sweeper_class'] = imex_1st_order 

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

90 description['level_params'] = level_params # pass level parameters 

91 description['step_params'] = step_params # pass step parameters 

92 description['space_transfer_class'] = fft_to_fft 

93 

94 # set time parameters 

95 t0 = 0.0 

96 Tend = 1.0 

97 

98 f = None 

99 if rank == 0: 

100 Path("data").mkdir(parents=True, exist_ok=True) 

101 f = open('data/step_7_B_out.txt', 'a') 

102 out = f'Running with ml={ml} and num_procs={num_procs}...' 

103 f.write(out + '\n') 

104 print(out) 

105 

106 # instantiate controller 

107 controller = controller_nonMPI(num_procs=num_procs, controller_params=controller_params, description=description) 

108 

109 # get initial values on finest level 

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

111 uinit = P.u_exact(t0) 

112 

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

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

115 uex = P.u_exact(Tend) 

116 err = abs(uex - uend) 

117 

118 if rank == 0: 

119 # filter statistics by type (number of iterations) 

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

121 

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

123 out = ( 

124 f' Min/Mean/Max number of iterations: ' 

125 f'{np.min(niters):4.2f} / {np.mean(niters):4.2f} / {np.max(niters):4.2f}' 

126 ) 

127 f.write(out + '\n') 

128 print(out) 

129 out = ' Range of values for number of iterations: %2i ' % np.ptp(niters) 

130 f.write(out + '\n') 

131 print(out) 

132 out = ' Position of max/min number of iterations: %2i -- %2i' % ( 

133 int(np.argmax(niters)), 

134 int(np.argmin(niters)), 

135 ) 

136 f.write(out + '\n') 

137 print(out) 

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

139 f.write(out + '\n') 

140 print(out) 

141 

142 out = f'Error: {err:6.4e}' 

143 f.write(out + '\n') 

144 print(out) 

145 

146 timing = get_sorted(stats, type='timing_run', sortby='time') 

147 out = f'Time to solution: {timing[0][1]:6.4f} sec.' 

148 f.write(out + '\n') 

149 print(out) 

150 

151 assert err <= 1.133e-05, 'Error is too high, got %s' % err 

152 if ml: 

153 if num_procs > 1: 

154 maxmean = 12.5 

155 else: 

156 maxmean = 6.6 

157 else: 

158 maxmean = 12.7 

159 assert np.mean(niters) <= maxmean, 'Mean number of iterations is too high, got %s' % np.mean(niters) 

160 

161 f.write('\n') 

162 print() 

163 f.close() 

164 

165 

166# %% [markdown] 

167# ## All runs 

168# 

169# SDC, MLSDC and PFASST with 10 steps in parallel (emulated in one process by `controller_nonMPI`), each in real 

170# and in spectral space. The code runs in serial with 

171# `python B_pySDC_with_mpi4pyfft.py`, and in parallel in space with, e.g., `mpirun -np 2 python 

172# B_pySDC_with_mpi4pyfft.py`. 

173 

174 

175# %% 

176def main(): 

177 """ 

178 Little helper routine to run the whole thing 

179 

180 Note: This can also be run with "mpirun -np 2 python B_pySDC_with_mpi4pyfft.py" 

181 """ 

182 run_simulation(spectral=False, ml=False, num_procs=1) 

183 run_simulation(spectral=True, ml=False, num_procs=1) 

184 run_simulation(spectral=False, ml=True, num_procs=1) 

185 run_simulation(spectral=True, ml=True, num_procs=1) 

186 run_simulation(spectral=False, ml=True, num_procs=10) 

187 run_simulation(spectral=True, ml=True, num_procs=10) 

188 

189 

190if __name__ == "__main__": 

191 main() 

192 

193# %% [markdown] 

194# ## Results 

195# 

196# mpi4py-fft does not run in the browser, nor in the environment this website is built in. These are the results of 

197# our CI, which runs this part in an environment with mpi4py-fft, in the run that built this page: 

198# 

199# :::{literalinclude} /../../data/step_7_B_out.txt 

200# :language: text 

201# ::: 

202# 

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

204# - The nonlinear Schrödinger example is not expected to work well with PFASST. In fact, SDC and MLSDC converge for 

205# larger time steps, but PFASST does not. 

206# :::