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

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.implementations.sweeper_classes.imex_1st_order import imex_1st_order 

8from pySDC.core.convergence_controller import ConvergenceController 

9from pySDC.implementations.convergence_controller_classes.estimate_extrapolation_error import ( 

10 EstimateExtrapolationErrorNonMPI, 

11) 

12from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence 

13 

14from pySDC.core.errors import ConvergenceError 

15 

16import numpy as np 

17 

18 

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

20 import pickle 

21 import os 

22 

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

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

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

26 data = pickle.load(file) 

27 else: 

28 from pySDC.helpers.stats_helper import get_sorted 

29 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity 

30 

31 convergence_controllers = { 

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

33 } 

34 desc = {'convergence_controllers': convergence_controllers} 

35 

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

37 t0 = 0 if t_init is None else t_init 

38 

39 if t == t0: 

40 return u0 

41 else: 

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

43 _t0 = t0 if _t0 is None else _t0 

44 

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

46 

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

48 data = u[1] 

49 

50 if t0 == 0: 

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

52 pickle.dump(data, file) 

53 

54 return data 

55 

56 

57RayleighBenard._u_exact = RayleighBenard.u_exact 

58RayleighBenard.u_exact = u_exact 

59 

60PROBLEM_PARAMS = {'Rayleigh': 2e4, 'nx': 256, 'nz': 128} 

61 

62 

63class ReachTendExactly(ConvergenceController): 

64 """ 

65 This convergence controller will adapt the step size of (hopefully) the last step such that `Tend` is reached very closely. 

66 Please pass the same `Tend` that you pass to the controller to the params for this to work. 

67 """ 

68 

69 def setup(self, controller, params, description, **kwargs): 

70 defaults = { 

71 "control_order": +94, 

72 "Tend": None, 

73 'min_step_size': 1e-10, 

74 } 

75 return {**defaults, **super().setup(controller, params, description, **kwargs)} 

76 

77 def get_new_step_size(self, controller, step, **kwargs): 

78 L = step.levels[0] 

79 

80 if not CheckConvergence.check_convergence(step): 

81 return None 

82 

83 dt = L.status.dt_new if L.status.dt_new else L.params.dt 

84 time_left = self.params.Tend - L.time - L.dt 

85 if time_left <= dt + self.params.min_step_size and not step.status.restart and time_left > 0: 

86 dt_new = ( 

87 min([dt + self.params.min_step_size, max([time_left, self.params.min_step_size])]) 

88 + np.finfo('float').eps * 10 

89 ) 

90 if dt_new != L.status.dt_new: 

91 L.status.dt_new = dt_new 

92 self.log( 

93 f'Reducing step size from {dt:12e} to {L.status.dt_new:.12e} because there is only {time_left:.12e} left.', 

94 step, 

95 ) 

96 

97 

98def run_RBC( 

99 custom_description=None, 

100 num_procs=1, 

101 Tend=21.0, 

102 hook_class=LogData, 

103 fault_stuff=None, 

104 custom_controller_params=None, 

105 u0=None, 

106 t0=20.0, 

107 use_MPI=False, 

108 **kwargs, 

109): 

110 """ 

111 Args: 

112 custom_description (dict): Overwrite presets 

113 num_procs (int): Number of steps for MSSDC 

114 Tend (float): Time to integrate to 

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

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

117 custom_controller_params (dict): Overwrite presets 

118 u0 (dtype_u): Initial value 

119 t0 (float): Starting time 

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

121 

122 Returns: 

123 dict: The stats object 

124 controller: The controller 

125 bool: If the code crashed 

126 """ 

127 EstimateExtrapolationErrorNonMPI.get_extrapolated_error = get_extrapolated_error_DAE 

128 

129 level_params = {} 

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

131 level_params['restol'] = -1 

132 

133 sweeper_params = {} 

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

135 sweeper_params['num_nodes'] = 3 

136 sweeper_params['QI'] = 'LU' 

137 sweeper_params['QE'] = 'PIC' 

138 

139 from mpi4py import MPI 

140 

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

142 

143 step_params = {} 

144 step_params['maxiter'] = 5 

145 

146 convergence_controllers = {} 

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

148 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeRounding 

149 

150 convergence_controllers[StepSizeRounding] = {} 

151 

152 controller_params = {} 

153 controller_params['logger_level'] = 30 

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

155 controller_params['mssdc_jac'] = False 

156 

157 if custom_controller_params is not None: 

158 controller_params = {**controller_params, **custom_controller_params} 

159 

160 imex_1st_order.compute_residual = compute_residual_DAE 

161 

162 description = {} 

163 description['problem_class'] = RayleighBenard 

164 description['problem_params'] = problem_params 

165 description['sweeper_class'] = imex_1st_order 

166 description['sweeper_params'] = sweeper_params 

167 description['level_params'] = level_params 

168 description['step_params'] = step_params 

169 description['convergence_controllers'] = convergence_controllers 

170 

171 if custom_description is not None: 

172 description = merge_descriptions(description, custom_description) 

173 

174 controller_args = { 

175 'controller_params': controller_params, 

176 'description': description, 

177 } 

178 if use_MPI: 

179 from mpi4py import MPI 

180 from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

181 

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

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

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

185 else: 

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

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

188 

189 t0 = 0.0 if t0 is None else t0 

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

191 

192 if fault_stuff is not None: 

193 from pySDC.projects.Resilience.fault_injection import prepare_controller_for_faults 

194 

195 prepare_controller_for_faults(controller, fault_stuff) 

196 

197 crash = False 

198 try: 

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

200 except ConvergenceError as e: 

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

202 stats = controller.return_stats() 

203 crash = True 

204 return stats, controller, crash 

205 

206 

207def generate_data_for_fault_stats(): 

208 prob = RayleighBenard(**PROBLEM_PARAMS) 

209 _ts = np.linspace(0, 22, 220, dtype=float) 

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

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

212 

213 

214def plot_order(t, dt, steps, num_nodes, e_tol=1e-9, restol=1e-9, ax=None, recompute=False): 

215 from pySDC.implementations.hooks.log_errors import LogGlobalErrorPostRun 

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

217 from pySDC.implementations.convergence_controller_classes.crash import StopAtNan 

218 from pySDC.helpers.stats_helper import get_sorted 

219 import pickle 

220 import os 

221 

222 sweeper_params = {'num_nodes': num_nodes} 

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

224 step_params = {'maxiter': 99} 

225 convergence_controllers = {StopAtNan: {}} 

226 

227 errors = [] 

228 dts = [] 

229 for i in range(steps): 

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

231 

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

233 

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

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

236 stats = pickle.load(file) 

237 else: 

238 

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

240 

241 desc = { 

242 'sweeper_params': sweeper_params, 

243 'level_params': level_params, 

244 'step_params': step_params, 

245 'convergence_controllers': convergence_controllers, 

246 } 

247 

248 stats, _, _ = run_RBC( 

249 Tend=t + dt, 

250 t0=t, 

251 custom_description=desc, 

252 hook_class=[LogGlobalErrorPostRun, LogSDCIterations, LogWork], 

253 ) 

254 

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

256 pickle.dump(stats, file) 

257 

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

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

260 

261 if len(e) > 0: 

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

263 else: 

264 errors += [np.nan] 

265 

266 errors = np.array(errors) 

267 dts = np.array(dts) 

268 mask = np.isfinite(errors) 

269 max_error = np.nanmax(errors) 

270 

271 errors = errors[mask] 

272 dts = dts[mask] 

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

274 ax.loglog( 

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

276 ) 

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

278 ax.set_ylabel(r'global error') 

279 ax.legend(frameon=False) 

280 

281 

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

283 prob = RayleighBenard(**PROBLEM_PARAMS) 

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

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

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

287 

288 import matplotlib.pyplot as plt 

289 

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

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

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

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

294 plt.show() 

295 

296 

297def plot_step_size(t0=0, Tend=30, e_tol=1e-3, recompute=False): 

298 import matplotlib.pyplot as plt 

299 import pickle 

300 import os 

301 from pySDC.helpers.stats_helper import get_sorted 

302 

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

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

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

306 stats = pickle.load(file) 

307 else: 

308 from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity 

309 

310 convergence_controllers = { 

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

312 } 

313 desc = {'convergence_controllers': convergence_controllers} 

314 

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

316 

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

318 pickle.dump(stats, file) 

319 

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

321 

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

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

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

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

326 ax.set_yscale('log') 

327 plt.show() 

328 

329 

330if __name__ == '__main__': 

331 # plot_step_size(0, 3) 

332 generate_data_for_fault_stats() 

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

334 # stats, _, _ = run_RBC()