Coverage for pySDC/projects/Resilience/RBC.py: 16%

114 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-04-01 13:12 +0000

1# script to run a Rayleigh-Benard Convection problem 

2from pySDC.implementations.problem_classes.generic_spectral import compute_residual_DAE, get_extrapolated_error_DAE 

3from pySDC.implementations.problem_classes.RayleighBenard import RayleighBenard 

4from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

5from pySDC.projects.Resilience.hook import hook_collection, LogData 

6from pySDC.projects.Resilience.strategies import merge_descriptions 

7from pySDC.projects.Resilience.sweepers import imex_1st_order_efficient 

8from pySDC.implementations.convergence_controller_classes.estimate_extrapolation_error import ( 

9 EstimateExtrapolationErrorNonMPI, 

10) 

11from pySDC.projects.Resilience.reachTendExactly import ReachTendExactly 

12 

13from pySDC.core.errors import ConvergenceError 

14 

15import numpy as np 

16 

17PROBLEM_PARAMS = {'Rayleigh': 3.2e5, 'nx': 256, 'nz': 128, 'max_cached_factorizations': 30} 

18 

19 

20def u_exact(self, t, u_init=None, t_init=None, recompute=False, _t0=None): 

21 import pickle 

22 import os 

23 

24 path = f'data/stats/RBC-u_init-{t:.8f}.pickle' 

25 if os.path.exists(path) and not recompute and t_init is None: 

26 with open(path, 'rb') as file: 

27 data = pickle.load(file) 

28 else: 

29 from pySDC.helpers.stats_helper import get_sorted 

30 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity 

31 

32 convergence_controllers = { 

33 Adaptivity: {'e_tol': 1e-8, 'dt_rel_min_slope': 0.25, 'dt_min': 1e-5}, 

34 } 

35 desc = {'convergence_controllers': convergence_controllers} 

36 

37 u0 = self._u_exact(0) if u_init is None else u_init 

38 t0 = 0 if t_init is None else t_init 

39 

40 if t == t0: 

41 return u0 

42 else: 

43 u0 = u0 if _t0 is None else self.u_exact(_t0) 

44 _t0 = t0 if _t0 is None else _t0 

45 

46 stats, _, _ = run_RBC(Tend=t, u0=u0, t0=_t0, custom_description=desc) 

47 

48 u = get_sorted(stats, type='u', recomputed=False)[-1] 

49 data = u[1] 

50 

51 if t0 == 0: 

52 with open(path, 'wb') as file: 

53 pickle.dump(data, file) 

54 

55 return data 

56 

57 

58if not hasattr(RayleighBenard, '_u_exact'): 

59 RayleighBenard._u_exact = RayleighBenard.u_exact 

60 RayleighBenard.u_exact = u_exact 

61 

62 

63def run_RBC( 

64 custom_description=None, 

65 num_procs=1, 

66 Tend=21.0, 

67 hook_class=LogData, 

68 fault_stuff=None, 

69 custom_controller_params=None, 

70 u0=None, 

71 t0=20.0, 

72 use_MPI=False, 

73 step_size_rounding=False, 

74 **kwargs, 

75): 

76 """ 

77 Args: 

78 custom_description (dict): Overwrite presets 

79 num_procs (int): Number of steps for MSSDC 

80 Tend (float): Time to integrate to 

81 hook_class (pySDC.Hook): A hook to store data 

82 fault_stuff (dict): A dictionary with information on how to add faults 

83 custom_controller_params (dict): Overwrite presets 

84 u0 (dtype_u): Initial value 

85 t0 (float): Starting time 

86 use_MPI (bool): Whether or not to use MPI 

87 

88 Returns: 

89 dict: The stats object 

90 controller: The controller 

91 bool: If the code crashed 

92 """ 

93 EstimateExtrapolationErrorNonMPI.get_extrapolated_error = get_extrapolated_error_DAE 

94 

95 level_params = {} 

96 level_params['dt'] = 1e-3 

97 level_params['restol'] = -1 

98 

99 sweeper_params = {} 

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

101 sweeper_params['num_nodes'] = 3 

102 sweeper_params['QI'] = 'LU' 

103 sweeper_params['QE'] = 'PIC' 

104 sweeper_params['initial_guess'] = 'copy' 

105 

106 from mpi4py import MPI 

107 

108 problem_params = {'comm': MPI.COMM_SELF, **PROBLEM_PARAMS} 

109 

110 step_params = {} 

111 step_params['maxiter'] = 5 

112 

113 convergence_controllers = {} 

114 convergence_controllers[ReachTendExactly] = {'Tend': Tend} 

115 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeRounding 

116 

117 if step_size_rounding: 

118 convergence_controllers[StepSizeRounding] = {} 

119 

120 controller_params = {} 

121 controller_params['logger_level'] = 15 

122 controller_params['hook_class'] = hook_collection + (hook_class if type(hook_class) == list else [hook_class]) 

123 controller_params['mssdc_jac'] = False 

124 

125 if custom_controller_params is not None: 

126 controller_params = {**controller_params, **custom_controller_params} 

127 

128 imex_1st_order_efficient.compute_residual = compute_residual_DAE 

129 

130 description = {} 

131 description['problem_class'] = RayleighBenard 

132 description['problem_params'] = problem_params 

133 description['sweeper_class'] = imex_1st_order_efficient 

134 description['sweeper_params'] = sweeper_params 

135 description['level_params'] = level_params 

136 description['step_params'] = step_params 

137 description['convergence_controllers'] = convergence_controllers 

138 

139 if custom_description is not None: 

140 description = merge_descriptions(description, custom_description) 

141 

142 controller_args = { 

143 'controller_params': controller_params, 

144 'description': description, 

145 } 

146 if use_MPI: 

147 from mpi4py import MPI 

148 from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

149 

150 comm = kwargs.get('comm', MPI.COMM_WORLD) 

151 controller = controller_MPI(**controller_args, comm=comm) 

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

153 else: 

154 controller = controller_nonMPI(**controller_args, num_procs=num_procs) 

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

156 

157 t0 = 0.0 if t0 is None else t0 

158 uinit = P.u_exact(t0) if u0 is None else u0 

159 

160 if fault_stuff is not None: 

161 from pySDC.projects.Resilience.fault_injection import prepare_controller_for_faults 

162 

163 prepare_controller_for_faults(controller, fault_stuff) 

164 

165 crash = False 

166 try: 

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

168 except ConvergenceError as e: 

169 print(f'Warning: Premature termination!: {e}') 

170 stats = controller.return_stats() 

171 crash = True 

172 return stats, controller, crash 

173 

174 

175def generate_data_for_fault_stats(Tend): 

176 prob = RayleighBenard(**PROBLEM_PARAMS) 

177 _ts = np.linspace(0, Tend, Tend * 10 + 1, dtype=float) 

178 for i in range(len(_ts) - 1): 

179 print(f'Generating reference solution from {_ts[i]:.4e} to {_ts[i+1]:.4e}') 

180 prob.u_exact(_ts[i + 1], _t0=_ts[i], recompute=False) 

181 

182 

183def plot_order(t, dt, steps, num_nodes, e_tol=1e-9, restol=1e-9, ax=None, recompute=False): # pragma: no cover 

184 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun 

185 from pySDC.implementations.hooks.log_work import LogSDCIterations, LogWork 

186 from pySDC.implementations.convergence_controller_classes.crash import StopAtNan 

187 from pySDC.helpers.stats_helper import get_sorted 

188 import pickle 

189 import os 

190 

191 sweeper_params = {'num_nodes': num_nodes} 

192 level_params = {'e_tol': e_tol, 'restol': restol} 

193 step_params = {'maxiter': 99} 

194 convergence_controllers = {StopAtNan: {}} 

195 

196 errors = [] 

197 dts = [] 

198 for i in range(steps): 

199 dts += [dt / 2**i] 

200 

201 path = f'data/stats/RBC-u_order-{t:.8f}-{dts[-1]:.8e}-{num_nodes}-{e_tol:.2e}-{restol:.2e}.pickle' 

202 

203 if os.path.exists(path) and not recompute: 

204 with open(path, 'rb') as file: 

205 stats = pickle.load(file) 

206 else: 

207 

208 level_params['dt'] = dts[-1] 

209 

210 desc = { 

211 'sweeper_params': sweeper_params, 

212 'level_params': level_params, 

213 'step_params': step_params, 

214 'convergence_controllers': convergence_controllers, 

215 } 

216 

217 stats, _, _ = run_RBC( 

218 Tend=t + dt, 

219 t0=t, 

220 custom_description=desc, 

221 hook_class=[LogGlobalErrorPostRun, LogSDCIterations, LogWork], 

222 ) 

223 

224 with open(path, 'wb') as file: 

225 pickle.dump(stats, file) 

226 

227 e = get_sorted(stats, type='e_global_post_run') 

228 # k = get_sorted(stats, type='k') 

229 

230 if len(e) > 0: 

231 errors += [e[-1][1]] 

232 else: 

233 errors += [np.nan] 

234 

235 errors = np.array(errors) 

236 dts = np.array(dts) 

237 mask = np.isfinite(errors) 

238 max_error = np.nanmax(errors) 

239 

240 errors = errors[mask] 

241 dts = dts[mask] 

242 ax.loglog(dts, errors, label=f'{num_nodes} nodes', marker='x') 

243 ax.loglog( 

244 dts, [max_error * (me / dts[0]) ** (2 * num_nodes - 1) for me in dts], ls='--', label=f'order {2*num_nodes-1}' 

245 ) 

246 ax.set_xlabel(r'$\Delta t$') 

247 ax.set_ylabel(r'global error') 

248 ax.legend(frameon=False) 

249 

250 

251def check_order(t=14, dt=1e-1, steps=6): 

252 prob = RayleighBenard(**PROBLEM_PARAMS) 

253 _ts = [0, t, t + dt] 

254 for i in range(len(_ts) - 1): 

255 prob.u_exact(_ts[i + 1], _t0=_ts[i]) 

256 

257 import matplotlib.pyplot as plt 

258 

259 fig, ax = plt.subplots(1, 1) 

260 for num_nodes in [1, 2, 3, 4]: 

261 plot_order(t=t, dt=dt, steps=steps, num_nodes=num_nodes, ax=ax, restol=1e-9) 

262 ax.set_title(f't={t:.2f}, dt={dt:.2f}') 

263 plt.show() 

264 

265 

266def plot_step_size(t0=0, Tend=30, e_tol=1e-3, recompute=False): # pragma: no cover 

267 import matplotlib.pyplot as plt 

268 import pickle 

269 import os 

270 from pySDC.helpers.stats_helper import get_sorted 

271 

272 path = f'data/stats/RBC-u-{t0:.8f}-{Tend:.8f}-{e_tol:.2e}.pickle' 

273 if os.path.exists(path) and not recompute: 

274 with open(path, 'rb') as file: 

275 stats = pickle.load(file) 

276 else: 

277 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity 

278 

279 convergence_controllers = { 

280 Adaptivity: {'e_tol': e_tol, 'dt_rel_min_slope': 0.25, 'dt_min': 1e-5}, 

281 } 

282 desc = {'convergence_controllers': convergence_controllers} 

283 

284 stats, _, _ = run_RBC(Tend=Tend, t0=t0, custom_description=desc) 

285 

286 with open(path, 'wb') as file: 

287 pickle.dump(stats, file) 

288 

289 fig, ax = plt.subplots(1, 1) 

290 

291 dt = get_sorted(stats, type='dt', recomputed=False) 

292 ax.plot([me[0] for me in dt], [me[1] for me in dt]) 

293 ax.set_ylabel(r'$\Delta t$') 

294 ax.set_xlabel(r'$t$') 

295 ax.set_yscale('log') 

296 plt.show() 

297 

298 

299def plot_factorizations_over_time(t0=0, Tend=50, e_tol=1e-3, recompute=False, adaptivity_mode='dt'): # pragma: no cover 

300 import matplotlib.pyplot as plt 

301 import pickle 

302 import os 

303 from pySDC.helpers.stats_helper import get_sorted 

304 from pySDC.implementations.hooks.log_work import LogWork 

305 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity, AdaptivityPolynomialError 

306 from pySDC.helpers.plot_helper import figsize_by_journal, setup_mpl 

307 from pySDC.projects.Resilience.paper_plots import savefig 

308 

309 setup_mpl() 

310 

311 fig, axs = plt.subplots(1, 2, figsize=figsize_by_journal('TUHH_thesis', 1.0, 0.4)) 

312 

313 if adaptivity_mode == 'dt': 

314 adaptivity = Adaptivity 

315 elif adaptivity_mode == 'dt_k': 

316 adaptivity = AdaptivityPolynomialError 

317 

318 dt_controllers = { 

319 'basic': { 

320 adaptivity: { 

321 'e_tol': e_tol, 

322 } 

323 }, 

324 'min slope': {adaptivity: {'e_tol': e_tol, 'beta': 0.5, 'dt_rel_min_slope': 2}}, 

325 'fixed': {}, 

326 # 'rounding': {adaptivity: {'e_tol': e_tol, 'beta': 0.5, 'dt_rel_min_slope': 2}, StepSizeRounding: {}}, 

327 } 

328 

329 for name, params in dt_controllers.items(): 

330 if adaptivity_mode == 'dt': 

331 path = f'data/stats/RBC-u-{t0:.8f}-{Tend:.8f}-{e_tol:.2e}-{name}.pickle' 

332 elif adaptivity_mode == 'dt_k': 

333 path = f'data/stats/RBC-u-{t0:.8f}-{Tend:.8f}-{e_tol:.2e}-{name}-dtk.pickle' 

334 

335 if os.path.exists(path) and not recompute: 

336 with open(path, 'rb') as file: 

337 stats = pickle.load(file) 

338 else: 

339 

340 convergence_controllers = { 

341 **params, 

342 } 

343 desc = {'convergence_controllers': convergence_controllers} 

344 

345 if name == 'fixed': 

346 if adaptivity_mode == 'dt': 

347 desc['level_params'] = {'dt': 2e-2} 

348 elif adaptivity_mode == 'dt_k': 

349 desc['level_params'] = {'dt': 2e-3} 

350 elif adaptivity_mode == 'dt_k': 

351 desc['level_params'] = {'restol': 1e-7} 

352 

353 stats, _, _ = run_RBC( 

354 Tend=Tend, t0=t0, custom_description=desc, hook_class=LogWork, step_size_rounding=False 

355 ) 

356 

357 with open(path, 'wb') as file: 

358 pickle.dump(stats, file) 

359 

360 factorizations = get_sorted(stats, type='work_factorizations') 

361 rhs_evals = get_sorted(stats, type='work_rhs') 

362 axs[0].plot([me[0] for me in factorizations], np.cumsum([me[1] for me in factorizations]), label=name) 

363 axs[1].plot([me[0] for me in rhs_evals], np.cumsum([me[1] for me in rhs_evals]), label=name) 

364 

365 axs[0].set_ylabel(r'matrix factorizations') 

366 axs[1].set_ylabel(r'right hand side evaluations') 

367 axs[0].set_xlabel(r'$t$') 

368 axs[1].set_xlabel(r'$t$') 

369 axs[0].set_yscale('log') 

370 axs[1].set_yscale('log') 

371 axs[0].legend(frameon=False) 

372 savefig(fig, f'RBC_step_size_controller_{adaptivity_mode}') 

373 

374 

375if __name__ == '__main__': 

376 # plot_step_size(0, 30) 

377 generate_data_for_fault_stats(Tend=30) 

378 # plot_factorizations_over_time(e_tol=1e-3, adaptivity_mode='dt') 

379 # plot_factorizations_over_time(recompute=False, e_tol=1e-5, adaptivity_mode='dt_k') 

380 # check_order(t=20, dt=1., steps=7) 

381 # stats, _, _ = run_RBC()