Coverage for pySDC/projects/parallelSDC/preconditioner_playground_MPI.py: 97%

1import os 

2import pickle 

3from collections import namedtuple 

4 

5import numpy as np 

6from mpi4py import MPI 

7import argparse 

8 

9import matplotlib as mpl 

10 

11mpl.use("Agg") 

12import pySDC.helpers.plot_helper as plt_helper 

13from pySDC.helpers.stats_helper import get_sorted 

14 

15from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

16from pySDC.implementations.problem_classes.AdvectionEquation_ND_FD import advectionNd 

17from pySDC.implementations.problem_classes.GeneralizedFisher_1D_FD_implicit import generalized_fisher 

18from pySDC.implementations.problem_classes.HeatEquation_ND_FD import heatNd_unforced 

19from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol 

20from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI 

21 

22ID = namedtuple('ID', ['setup', 'qd_type', 'param']) 

23 

24 

25def main(comm=None): 

26 # initialize level parameters (part I) 

27 level_params = dict() 

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

29 

30 # initialize sweeper parameters (part I) 

31 sweeper_params = dict() 

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

33 sweeper_params['num_nodes'] = comm.Get_size() 

34 sweeper_params['comm'] = comm 

35 

36 # initialize step parameters 

37 step_params = dict() 

38 step_params['maxiter'] = 100 

39 

40 # initialize controller parameters 

41 controller_params = dict() 

42 controller_params['logger_level'] = 30 

43 

44 # set up list of Q-delta types and setups 

45 qd_list = ['IEpar', 'Qpar', 'MIN', 'MIN3', 'MIN_GT'] 

46 setup_list = [ 

47 ('heat', 63, [10.0**i for i in range(-3, 3)]), 

48 ('advection', 64, [10.0**i for i in range(-3, 3)]), 

49 ('vanderpol', 2, [0.1 * 2**i for i in range(0, 10)]), 

50 ('fisher', 63, [2**i for i in range(-2, 3)]), 

51 ] 

52 # setup_list = [('fisher', 63, [2 * i for i in range(1, 6)])] 

53 

54 # pre-fill results with lists of setups 

55 results = dict() 

56 for setup, nvars, param_list in setup_list: 

57 results[setup] = (nvars, param_list) 

58 

59 # loop over all Q-delta matrix types 

60 for qd_type in qd_list: 

61 # assign implicit Q-delta matrix 

62 sweeper_params['QI'] = qd_type 

63 

64 # loop over all setups 

65 for setup, nvars, param_list in setup_list: 

66 # initialize problem parameters (part I) 

67 problem_params = dict() 

68 if setup != 'vanderpol': 

69 problem_params['nvars'] = nvars # number of degrees of freedom for each level 

70 

71 # loop over all parameters 

72 for param in param_list: 

73 # fill description for the controller 

74 description = dict() 

75 description['sweeper_class'] = generic_implicit_MPI # pass sweeper 

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

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

78 # description['base_transfer_class'] = base_transfer_mpi 

79 

80 print('working on: %s - %s - %s' % (qd_type, setup, param)) 

81 

82 # decide which setup to take 

83 if setup == 'heat': 

84 problem_params['nu'] = param 

85 problem_params['freq'] = 2 

86 problem_params['bc'] = 'dirichlet-zero' # boundary conditions 

87 

88 level_params['dt'] = 0.1 

89 

90 description['problem_class'] = heatNd_unforced 

91 description['problem_params'] = problem_params 

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

93 

94 elif setup == 'advection': 

95 problem_params['c'] = param 

96 problem_params['order'] = 2 

97 problem_params['freq'] = 2 

98 problem_params['stencil_type'] = 'center' # boundary conditions 

99 problem_params['bc'] = 'periodic' # boundary conditions 

100 

101 level_params['dt'] = 0.1 

102 

103 description['problem_class'] = advectionNd 

104 description['problem_params'] = problem_params 

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

106 

107 elif setup == 'vanderpol': 

108 problem_params['newton_tol'] = 1e-09 

109 problem_params['newton_maxiter'] = 20 

110 problem_params['mu'] = param 

111 problem_params['u0'] = np.array([2.0, 0]) 

112 

113 level_params['dt'] = 0.1 

114 

115 description['problem_class'] = vanderpol 

116 description['problem_params'] = problem_params 

117 description['level_params'] = level_params 

118 

119 elif setup == 'fisher': 

120 problem_params['nu'] = 1 

121 problem_params['lambda0'] = param 

122 problem_params['newton_maxiter'] = 20 

123 problem_params['newton_tol'] = 1e-10 

124 problem_params['interval'] = (-5, 5) 

125 

126 level_params['dt'] = 0.01 

127 

128 description['problem_class'] = generalized_fisher 

129 description['problem_params'] = problem_params 

130 description['level_params'] = level_params 

131 

132 else: 

133 print('Setup not implemented..', setup) 

134 exit() 

135 

136 # instantiate controller 

137 controller = controller_nonMPI( 

138 num_procs=1, controller_params=controller_params, description=description 

139 ) 

140 

141 # get initial values on finest level 

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

143 uinit = P.u_exact(0) 

144 

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

146 uend, stats = controller.run(u0=uinit, t0=0, Tend=level_params['dt']) 

147 

148 # filter statistics by type (number of iterations) 

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

150 

151 # just one time-step, grep number of iteration and store 

152 niter = iter_counts[0][1] 

153 id = ID(setup=setup, qd_type=qd_type, param=param) 

154 results[id] = niter 

155 

156 assert len(results) == (6 + 6 + 10 + 5) * 5 + 4, 'ERROR: did not get all results, got %s' % len(results) 

157 

158 if comm.Get_rank() == 0: 

159 # write out for later visualization 

160 with open('data/parallelSDC_iterations_precond_MPI.pkl', 'wb') as file: 

161 pickle.dump(results, file) 

162 

163 assert os.path.isfile('data/parallelSDC_iterations_precond_MPI.pkl'), 'ERROR: pickle did not create file' 

164 

165 

166# pragma: no cover 

167def plot_iterations(): 

168 """ 

169 Helper routine to plot iteration counts 

170 """ 

171 

172 with open('data/parallelSDC_iterations_precond_MPI.pkl', 'rb') as file: 

173 results = pickle.load(file) 

174 

175 # find the lists/header required for plotting 

176 qd_type_list = [] 

177 setup_list = [] 

178 for key in results.keys(): 

179 if isinstance(key, ID): 

180 if key.qd_type not in qd_type_list: 

181 qd_type_list.append(key.qd_type) 

182 elif isinstance(key, str): 

183 setup_list.append(key) 

184 print('Found these type of preconditioners:', qd_type_list) 

185 print('Found these setups:', setup_list) 

186 

187 assert len(qd_type_list) == 5, 'ERROR did not find five preconditioners, got %s' % qd_type_list 

188 assert len(setup_list) == 4, 'ERROR: did not find four setup, got %s' % setup_list 

189 

190 plt_helper.setup_mpl() 

191 print('post setup') 

192 

193 # loop over setups and Q-delta types: one figure per setup, all Qds in one plot 

194 for setup in list(setup_list[:]): 

195 plot_setup(results, setup) 

196 mpl.pyplot.close("all") 

197 return 0 

198 

199 

200def plot_setup(results, setup): 

201 print(f'setup of {setup}') 

202 

203 qd_type_list = ['IEpar', 'Qpar', 'MIN', 'MIN3', 'MIN_GT'] 

204 marker_list = ['s', 'o', '^', 'v', 'x'] 

205 color_list = ['r', 'g', 'b', 'c', 'm'] 

206 

207 fig, ax = plt_helper.newfig(textwidth=238.96, scale=0.89) 

208 

209 for qd_type, marker, color in zip(qd_type_list, marker_list, color_list): 

210 niter = np.zeros(len(results[setup][1])) 

211 for key in results.keys(): 

212 if isinstance(key, ID): 

213 if key.setup == setup and key.qd_type == qd_type: 

214 xvalue = results[setup][1].index(key.param) 

215 niter[xvalue] = results[key] 

216 ls = '-' 

217 lw = 1 

218 ax.semilogx( 

219 results[setup][1], 

220 niter, 

221 label=qd_type, 

222 lw=lw, 

223 linestyle=ls, 

224 color=color, 

225 marker=marker, 

226 markeredgecolor='k', 

227 ) 

228 

229 if setup == 'heat': 

230 xlabel = r'$\nu$' 

231 elif setup == 'advection': 

232 xlabel = r'$c$' 

233 elif setup == 'fisher': 

234 xlabel = r'$\lambda_0$' 

235 elif setup == 'vanderpol': 

236 xlabel = r'$\mu$' 

237 else: 

238 print('Setup not implemented..', setup) 

239 exit() 

240 

241 ax.set_ylim(0, 60) 

242 fig.legend(loc=2, ncol=1) 

243 ax.set_ylabel('number of iterations') 

244 ax.set_xlabel(xlabel) 

245 ax.grid() 

246 

247 # save plot as PDF and PGF 

248 fname = 'data/parallelSDC_preconditioner_MPI_' + setup 

249 plt_helper.savefig(fname) 

250 del fig, ax 

251 

252 assert os.path.isfile(fname + '.pdf'), 'ERROR: plotting did not create PDF file' 

253 # assert os.path.isfile(fname + '.pgf'), 'ERROR: plotting did not create PGF file' 

254 assert os.path.isfile(fname + '.png'), 'ERROR: plotting did not create PNG file' 

255 print(f"Successfully stored image {fname}.png") 

256 

257 return 

258 

259 

260if __name__ == "__main__": 

261 parser = argparse.ArgumentParser() 

262 parser.add_argument("action", type=str, default="", const="", nargs='?') 

263 args = parser.parse_args() 

264 if args.action == "": 

265 comm = MPI.COMM_WORLD 

266 main(comm=comm) 

267 if comm.Get_rank() == 0: 

268 plot_iterations() 

269 elif args.action == "simulate": 

270 comm = MPI.COMM_WORLD 

271 main(comm=comm) 

272 elif args.action == "plot": 

273 plot_iterations() 

274 else: 

275 raise KeyError(f"Unknown action '{args.action}'. Known actions are 'simulate', 'plot', or none for both")