Coverage for pySDC/tutorial/step_7/C_pySDC_with_PETSc.py: 98%

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 C: Time-parallel pySDC with space-parallel PETSc 

14# 

15# With rather unfavorable scaling properties, parallel-in-time methods are only really useful when spatial 

16# parallelization is maxed out. To work with spatial parallelization, this part shows how to (1) include and work 

17# with an external library and (2) set up space- and time-parallel runs. We use the forced heat equation again as our 

18# testbed, and [PETSc](https://petsc.org/) for the space-parallel data structures and the linear solver. See 

19# `implementations/datatype_classes/petsc_dmda_grid.py` and 

20# `implementations/problem_classes/HeatEquation_2D_PETSc_forced.py` for the details of the PETSc bindings. 

21# 

22# ## Processes in space and in time 

23# 

24# The world communicator is split twice, by coloring: into space communicators of `num_procs_space` ranks each, which 

25# go to the problem class, and across them into time communicators, which go to `controller_MPI`. With 4 ranks and 

26# `num_procs_space = 2`, ranks 0 and 1 share space, as do ranks 2 and 3, and ranks 0 and 2 form one time communicator. 

27 

28# %% 

29import sys 

30from pathlib import Path 

31 

32import numpy as np 

33from mpi4py import MPI 

34 

35from pySDC.helpers.stats_helper import get_sorted 

36 

37from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

38from pySDC.implementations.problem_classes.HeatEquation_2D_PETSc_forced import heat2d_petsc_forced 

39from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order 

40from pySDC.implementations.transfer_classes.TransferPETScDMDA import mesh_to_mesh_petsc_dmda 

41 

42 

43def main(num_procs_space=None, fname='step_7_C_out.txt'): 

44 """ 

45 Program to demonstrate usage of PETSc data structures and spatial parallelization, 

46 combined with parallelization in time. 

47 

48 Args: 

49 num_procs_space (int): number of ranks to give the space communicator. Defaults to the 

50 first command line argument when run as a script, and to 1 otherwise. 

51 fname (str): name of the output file written under ``data/``. 

52 """ 

53 # set MPI communicator 

54 comm = MPI.COMM_WORLD 

55 

56 world_rank = comm.Get_rank() 

57 world_size = comm.Get_size() 

58 

59 # split world communicator to create space-communicators 

60 if num_procs_space is not None: 

61 color = int(world_rank / num_procs_space) 

62 else: 

63 color = int(world_rank / 1) 

64 space_comm = comm.Split(color=color) 

65 space_rank = space_comm.Get_rank() 

66 

67 # split world communicator to create time-communicators 

68 if num_procs_space is not None: 

69 color = int(world_rank % num_procs_space) 

70 else: 

71 color = int(world_rank / world_size) 

72 time_comm = comm.Split(color=color) 

73 time_rank = time_comm.Get_rank() 

74 

75 # initialize level parameters 

76 level_params = dict() 

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

78 level_params['dt'] = 0.125 

79 level_params['nsweeps'] = [1] 

80 

81 # initialize sweeper parameters 

82 sweeper_params = dict() 

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

84 sweeper_params['num_nodes'] = [3] 

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

86 sweeper_params['initial_guess'] = 'zero' 

87 

88 # initialize problem parameters 

89 problem_params = dict() 

90 problem_params['nu'] = 1.0 # diffusion coefficient 

91 problem_params['freq'] = 2 # frequency for the test value 

92 problem_params['cnvars'] = [(65, 65)] # number of degrees of freedom for the coarsest level 

93 problem_params['refine'] = [1, 0] # number of refinements 

94 problem_params['comm'] = space_comm # pass space-communicator to problem class 

95 problem_params['sol_tol'] = 1e-12 # set tolerance to PETSc' linear solver 

96 

97 # initialize step parameters 

98 step_params = dict() 

99 step_params['maxiter'] = 50 

100 

101 # initialize space transfer parameters 

102 space_transfer_params = dict() 

103 space_transfer_params['rorder'] = 2 

104 space_transfer_params['iorder'] = 2 

105 space_transfer_params['periodic'] = False 

106 

107 # initialize controller parameters 

108 controller_params = dict() 

109 controller_params['logger_level'] = 20 if space_rank == 0 else 99 # set level depending on rank 

110 controller_params['dump_setup'] = False 

111 

112 # fill description dictionary for easy step instantiation 

113 description = dict() 

114 description['problem_class'] = heat2d_petsc_forced # pass problem class 

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

116 description['sweeper_class'] = imex_1st_order # pass sweeper (see part B) 

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

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

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

120 description['space_transfer_class'] = mesh_to_mesh_petsc_dmda # pass spatial transfer class 

121 description['space_transfer_params'] = space_transfer_params # pass parameters for spatial transfer 

122 

123 # set time parameters 

124 t0 = 0.0 

125 Tend = 0.25 

126 

127 # instantiate controller 

128 controller = controller_MPI(controller_params=controller_params, description=description, comm=time_comm) 

129 

130 # get initial values on finest level 

131 P = controller.S.levels[0].prob 

132 uinit = P.u_exact(t0) 

133 

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

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

136 

137 # compute exact solution and compare 

138 uex = P.u_exact(Tend) 

139 err = abs(uex - uend) 

140 

141 # filter statistics by type (number of iterations) 

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

143 

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

145 

146 # limit output to space-rank 0 (as before when setting the logger level) 

147 if space_rank == 0: 

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

149 f = open('data/' + fname, 'a+') 

150 

151 out = 'This is time-rank %i...' % time_rank 

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

153 print(out) 

154 

155 # compute and print statistics 

156 for item in iter_counts: 

157 out = 'Number of iterations for time %4.2f: %2i' % item 

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

159 print(out) 

160 

161 out = ' Mean number of iterations: %4.2f' % np.mean(niters) 

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

163 print(out) 

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

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

166 print(out) 

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

168 int(np.argmax(niters)), 

169 int(np.argmin(niters)), 

170 ) 

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

172 print(out) 

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

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

175 print(out) 

176 

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

178 

179 out = 'Time to solution: %6.4f sec.' % timing[0][1] 

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

181 print(out) 

182 out = 'Error vs. PDE solution: %6.4e' % err 

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

184 print(out) 

185 

186 f.close() 

187 

188 assert err < 2e-04, 'ERROR: did not match error tolerance, got %s' % err 

189 assert np.mean(niters) <= 12, 'ERROR: number of iterations is too high, got %s' % np.mean(niters) 

190 

191 space_comm.Free() 

192 time_comm.Free() 

193 

194 

195# %% [markdown] 

196# ## Running it 

197# 

198# With `mpirun`, giving the number of ranks in space and the output file, e.g. for 2 in time and 2 in space: 

199# 

200# ```bash 

201# mpirun -np 4 python C_pySDC_with_PETSc.py 2 step_7_C_out_2x2.txt 

202# ``` 

203 

204# %% 

205if __name__ == "__main__": 

206 # still runnable straight from the shell, as the tutorial text describes: 

207 # mpirun -np 4 python C_pySDC_with_PETSc.py 2 step_7_C_out_2x2.txt 

208 main( 

209 num_procs_space=int(sys.argv[1]) if len(sys.argv) >= 2 else None, 

210 fname=sys.argv[2] if len(sys.argv) == 3 else 'step_7_C_out.txt', 

211 ) 

212 

213# %% [markdown] 

214# ## Results 

215# 

216# PETSc does not run in the browser, nor in the environment this website is built in. Our CI runs this part three 

217# times, in an environment with PETSc: with 1 and with 2 processes in space, and with 4 processes, 2 in time and 2 in 

218# space. Do not expect scaling, due to the CI environment. These are the results of the run that built this page, 

219# with 1 process: 

220# 

221# :::{literalinclude} /../../data/step_7_C_out_1x1.txt 

222# :language: text 

223# ::: 

224# 

225# With 2 processes in space: 

226# 

227# :::{literalinclude} /../../data/step_7_C_out_1x2.txt 

228# :language: text 

229# ::: 

230# 

231# And with 2 in time and 2 in space: 

232# 

233# :::{literalinclude} /../../data/step_7_C_out_2x2.txt 

234# :language: text 

235# ::: 

236# 

237# The iterations and the error agree in all three: distributing the spatial problem over processes changes how it is 

238# solved, not what, and running the steps in parallel does not change what PFASST computes. Only the time to solution 

239# differs, and each time rank prints its own steps. 

240# 

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

242# - We need processors in space and time, which can be achieved by `comm.Split` and coloring. The space communicator 

243# is then passed to the problem class. 

244# :::