Coverage for pySDC/projects/RayleighBenard/RBC3D_configs.py: 77%

299 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-12 05:46 +0000

1from pySDC.projects.GPU.configs.base_config import Config 

2 

3 

4def get_config(args): 

5 name = args['config'] 

6 if name == 'RBC3D': 

7 return RayleighBenard3DRegular(args) 

8 elif name in globals().keys(): 

9 return globals()[name](args) 

10 else: 

11 raise NotImplementedError(f'There is no configuration called {name!r}!') 

12 

13 

14class RayleighBenard3DRegular(Config): 

15 sweeper_type = 'IMEX' 

16 Tend = 50 

17 gamma = 1 

18 res_ratio = 1 

19 dealiasing = 3.0 / 2.0 

20 

21 def get_sweeper(self, useMPI): 

22 if useMPI: 

23 from pySDC.projects.RayleighBenard.sweepers import imex_1st_order_MPI_fixed_k as sweeper 

24 else: 

25 from pySDC.projects.RayleighBenard.sweepers import imex_1st_order_diagonal_serial as sweeper 

26 return sweeper 

27 

28 def get_file_name(self): 

29 res = self.args['res'] 

30 return f'{self.base_path}/data/{type(self).__name__}-res{res}.pySDC' 

31 

32 def get_LogToFile(self, *args, **kwargs): 

33 if self.comms[1].rank > 0: 

34 return None 

35 import numpy as np 

36 from pySDC.implementations.hooks.log_solution import LogToFile 

37 

38 LogToFile.filename = self.get_file_name() 

39 LogToFile.time_increment = 5e-1 

40 # LogToFile.allow_overwriting = True 

41 

42 return LogToFile 

43 

44 def get_controller_params(self, *args, **kwargs): 

45 from pySDC.implementations.hooks.log_step_size import LogStepSize 

46 

47 controller_params = super().get_controller_params(*args, **kwargs) 

48 controller_params['hook_class'] += [LogStepSize] 

49 return controller_params 

50 

51 def get_description(self, *args, MPIsweeper=False, res=-1, **kwargs): 

52 from pySDC.implementations.problem_classes.RayleighBenard3D import ( 

53 RayleighBenard3D, 

54 ) 

55 from pySDC.implementations.problem_classes.generic_spectral import ( 

56 compute_residual_DAE, 

57 compute_residual_DAE_MPI, 

58 ) 

59 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeSlopeLimiter 

60 from pySDC.implementations.convergence_controller_classes.crash import StopAtNan 

61 

62 desc = super().get_description(*args, MPIsweeper=MPIsweeper, **kwargs) 

63 

64 if MPIsweeper: 

65 desc['sweeper_class'].compute_residual = compute_residual_DAE_MPI 

66 else: 

67 desc['sweeper_class'].compute_residual = compute_residual_DAE 

68 

69 desc['level_params']['dt'] = 0.01 

70 desc['level_params']['restol'] = 1e-7 

71 

72 desc['convergence_controllers'][StepSizeSlopeLimiter] = {'dt_rel_min_slope': 0.1} 

73 desc['convergence_controllers'][StopAtNan] = {} 

74 

75 desc['sweeper_params']['quad_type'] = 'RADAU-RIGHT' 

76 desc['sweeper_params']['num_nodes'] = 2 

77 desc['sweeper_params']['QI'] = 'MIN-SR-S' 

78 desc['sweeper_params']['QE'] = 'PIC' 

79 

80 res = 64 if res == -1 else res 

81 desc['problem_params']['Rayleigh'] = 1e8 

82 desc['problem_params']['nx'] = self.res_ratio * res 

83 desc['problem_params']['ny'] = self.res_ratio * res 

84 desc['problem_params']['nz'] = res 

85 desc['problem_params']['Lx'] = self.gamma 

86 desc['problem_params']['Ly'] = self.gamma 

87 desc['problem_params']['Lz'] = 1 

88 desc['problem_params']['heterogeneous'] = True 

89 desc['problem_params']['dealiasing'] = self.dealiasing 

90 

91 desc['step_params']['maxiter'] = 3 

92 

93 desc['problem_class'] = RayleighBenard3D 

94 

95 return desc 

96 

97 def get_initial_condition(self, P, *args, restart_idx=0, **kwargs): 

98 

99 if restart_idx == 0: 

100 u0 = P.u_exact(t=0, seed=P.comm.rank, noise_level=1e-3) 

101 u0_with_pressure = P.solve_system(u0, 1e-9, u0) 

102 P.cached_factorizations.pop(1e-9) 

103 return u0_with_pressure, 0 

104 else: 

105 from pySDC.helpers.fieldsIO import FieldsIO 

106 

107 P.setUpFieldsIO() 

108 outfile = FieldsIO.fromFile(self.get_file_name()) 

109 

110 t0, solution = outfile.readField(restart_idx) 

111 solution = solution[: P.spectral.ncomponents, ...] 

112 

113 if P.useGPU: 

114 solution = P.xp.array(solution) 

115 

116 u0 = P.u_init 

117 

118 if P.spectral_space: 

119 u0[...] = P.transform(solution) 

120 else: 

121 u0[...] = solution 

122 

123 return u0, t0 

124 

125 def prepare_caches(self, prob): 

126 """ 

127 Cache the fft objects, which are expensive to create on GPU because graphs have to be initialized. 

128 """ 

129 prob.eval_f(prob.u_init) 

130 

131 def prepare_for_benchmark(self): 

132 def _pass(*args, **kwargs): 

133 pass 

134 

135 self.get_LogToFile = _pass 

136 

137 def prepare_caches_for_benchmark(self, prob, controller): 

138 _rhs = prob.u_init 

139 sweeper = controller.MS[0].levels[0].sweep 

140 _dt = sweeper.level.dt 

141 

142 hooks = controller.hooks 

143 

144 # mute controller 

145 type(controller).hooks = [] 

146 controller.run(_rhs, 0, _dt) 

147 

148 # unmute controller 

149 controller.hooks = hooks 

150 

151 try: 

152 import cuda as cp 

153 

154 cp.cuda.get_current_stream().synchronize() 

155 except ModuleNotFoundError: 

156 pass 

157 from mpi4py import MPI 

158 

159 MPI.COMM_WORLD.Barrier() 

160 controller.logger.critical('Set up caches for benchmarking') 

161 

162 def prepare_description_for_benchmark(self, description, controller_params): 

163 from pySDC.projects.RayleighBenard.benchmarks.print_timings_hook import PrintCPUTimings, PrintGPUTimings 

164 from mpi4py import MPI 

165 

166 self.Tend = 5 * description['level_params']['dt'] 

167 

168 controller_params['logger_level'] = 40 

169 controller_params['hook_class'] += [PrintCPUTimings] 

170 if description['problem_params']['useGPU']: 

171 controller_params['hook_class'] += [PrintGPUTimings] 

172 

173 description['problem_params']['max_cached_factorizations'] = 99 

174 

175 time_rank = 0 

176 if 'comm' in description['sweeper_params'].keys(): 

177 time_rank = description['sweeper_params']['comm'].rank 

178 for i in range(MPI.COMM_WORLD.size): 

179 if MPI.COMM_WORLD.rank == i: 

180 print( 

181 f'Global rank {MPI.COMM_WORLD.rank} is {time_rank} in time and {description["problem_params"]["comm"].rank} in space', 

182 flush=True, 

183 ) 

184 MPI.COMM_WORLD.barrier() 

185 

186 

187class RBC3Dverification(RayleighBenard3DRegular): 

188 converged = 0 

189 dt = 1e-2 

190 ic_config = { 

191 'config': None, 

192 'res': -1, 

193 'dt': -1, 

194 } 

195 res = None 

196 Ra = None 

197 Tend = 100 

198 res_ratio = 4 

199 gamma = 4 

200 

201 def get_file_name(self): 

202 res = self.args['res'] 

203 dt = self.args['dt'] 

204 return f'{self.base_path}/data/{type(self).__name__}-res{res}-dt{dt:.0e}.pySDC' 

205 

206 def get_description(self, *args, res=-1, dt=-1, **kwargs): 

207 desc = super().get_description(*args, **kwargs) 

208 desc['level_params']['nsweeps'] = 4 

209 desc['level_params']['restol'] = -1 

210 desc['step_params']['maxiter'] = 1 

211 desc['sweeper_params']['QI'] = 'MIN-SR-S' 

212 desc['sweeper_params']['skip_residual_computation'] = ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE') 

213 desc['sweeper_params']['num_nodes'] = 4 

214 Ra = int(type(self).__name__[-3]) * 10 ** int(type(self).__name__[-1]) 

215 desc['problem_params']['Rayleigh'] = Ra 

216 desc['problem_params']['Prandtl'] = 0.7 

217 

218 _res = self.res if res == -1 else res 

219 desc['problem_params']['nx'] = _res * self.res_ratio 

220 desc['problem_params']['ny'] = _res * self.res_ratio 

221 desc['problem_params']['nz'] = _res 

222 

223 _dt = self.dt if dt == -1 else dt 

224 desc['level_params']['dt'] = _dt 

225 

226 desc['problem_params']['Lx'] = float(self.gamma) 

227 desc['problem_params']['Ly'] = float(self.gamma) 

228 desc['problem_params']['Lz'] = 1.0 

229 return desc 

230 

231 def get_initial_condition(self, P, *args, restart_idx=0, **kwargs): 

232 if self.ic_config['config'] is None or restart_idx != 0: 

233 return super().get_initial_condition(P, *args, restart_idx=restart_idx, **kwargs) 

234 

235 # read initial conditions 

236 from pySDC.helpers.fieldsIO import FieldsIO 

237 

238 ic_config = self.ic_config['config']( 

239 args={**self.args, 'res': self.ic_config['res'], 'dt': self.ic_config['dt']} 

240 ) 

241 ic_config.base_path = self.base_path 

242 desc = ic_config.get_description(res=self.ic_config['res'], dt=self.ic_config['dt']) 

243 ic_nx = desc['problem_params']['nx'] 

244 ic_ny = desc['problem_params']['ny'] 

245 ic_nz = desc['problem_params']['nz'] 

246 

247 _P = type(P)( 

248 nx=ic_nx, 

249 ny=ic_ny, 

250 nz=ic_nz, 

251 comm=P.comm, 

252 useGPU=P.useGPU, 

253 Dirichlet_recombination=False, 

254 left_preconditioner=False, 

255 ) 

256 _P.setUpFieldsIO() 

257 filename = ic_config.get_file_name() 

258 ic_file = FieldsIO.fromFile(filename) 

259 t0, ics = ic_file.readField(-1) 

260 ics = ics[: P.spectral.ncomponents, ...] 

261 P.logger.info(f'Loaded initial conditions from {filename!r} at t={t0}.') 

262 

263 # interpolate the initial conditions using padded transforms 

264 padding = (P.nx / ic_nx, P.ny / ic_ny, P.nz / ic_nz) 

265 P.logger.info(f'Interpolating initial conditions from {ic_nx}x{ic_ny}x{ic_nz} to {P.nx}x{P.ny}x{P.nz}') 

266 

267 ics = _P.xp.array(ics) 

268 _ics_hat = _P.transform(ics) 

269 ics_interpolated = _P.itransform(_ics_hat, padding=padding) 

270 

271 self.get_LogToFile() 

272 

273 P.setUpFieldsIO() 

274 if P.spectral_space: 

275 u0_hat = P.u_init_forward 

276 u0_hat[...] = P.transform(ics_interpolated) 

277 return u0_hat, 0 

278 else: 

279 return ics_interpolated, 0 

280 

281 

282class RBC3DM2K3(RBC3Dverification): 

283 

284 def get_description(self, *args, **kwargs): 

285 desc = super().get_description(*args, **kwargs) 

286 desc['level_params']['nsweeps'] = 3 

287 desc['sweeper_params']['num_nodes'] = 2 

288 return desc 

289 

290 

291class RBC3DM2K2(RBC3Dverification): 

292 

293 def get_description(self, *args, **kwargs): 

294 desc = super().get_description(*args, **kwargs) 

295 desc['level_params']['nsweeps'] = 2 

296 desc['sweeper_params']['num_nodes'] = 2 

297 return desc 

298 

299 

300class RBC3DM3K4(RBC3Dverification): 

301 

302 def get_description(self, *args, **kwargs): 

303 desc = super().get_description(*args, **kwargs) 

304 desc['level_params']['nsweeps'] = 4 

305 desc['sweeper_params']['num_nodes'] = 3 

306 return desc 

307 

308 

309class RBC3DM4K4(RBC3Dverification): 

310 

311 def get_description(self, *args, **kwargs): 

312 desc = super().get_description(*args, **kwargs) 

313 desc['level_params']['nsweeps'] = 4 

314 desc['sweeper_params']['num_nodes'] = 4 

315 return desc 

316 

317 

318class RBC3DverificationRK(RBC3Dverification): 

319 

320 def get_description(self, *args, res=-1, dt=-1, **kwargs): 

321 from pySDC.implementations.sweeper_classes.Runge_Kutta import ARK3 

322 

323 desc = super().get_description(*args, res=res, dt=dt, **kwargs) 

324 desc['level_params']['nsweeps'] = 1 

325 desc['level_params']['restol'] = -1 

326 desc['step_params']['maxiter'] = 1 

327 desc['sweeper_params']['skip_residual_computation'] = ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE') 

328 desc['sweeper_params'].pop('QI') 

329 desc['sweeper_params'].pop('num_nodes') 

330 desc['sweeper_class'] = ARK3 

331 return desc 

332 

333 

334class RBC3DverificationEuler(RBC3DverificationRK): 

335 

336 def get_description(self, *args, res=-1, dt=-1, **kwargs): 

337 from pySDC.implementations.sweeper_classes.Runge_Kutta import IMEXEulerStifflyAccurate 

338 

339 desc = super().get_description(*args, res=res, dt=dt, **kwargs) 

340 desc['sweeper_class'] = IMEXEulerStifflyAccurate 

341 return desc 

342 

343 

344# --- Ra 1e5 --- 

345class RBC3DG4R4SDC22Ra1e5(RBC3DM2K2): 

346 Tend = 200 

347 dt = 4e-2 

348 res = 32 

349 converged = 50 

350 

351 

352class RBC3DG4R4SDC23Ra1e5(RBC3DM2K3): 

353 Tend = 200 

354 dt = 6e-2 

355 res = 32 

356 converged = 50 

357 

358 

359class RBC3DG4R4SDC34Ra1e5(RBC3DM3K4): 

360 Tend = 200 

361 dt = 6e-2 

362 res = 32 

363 converged = 50 

364 

365 

366class RBC3DG4R4SDC44Ra1e5(RBC3DM4K4): 

367 Tend = 200 

368 dt = 6e-2 

369 res = 32 

370 converged = 50 

371 

372 

373class RBC3DG4R4RKRa1e5(RBC3DverificationRK): 

374 Tend = 200 

375 dt = 5e-2 

376 res = 32 

377 converged = 50 

378 

379 

380class RBC3DG4R4EulerRa1e5(RBC3DverificationEuler): 

381 Tend = 200 

382 dt = 2e-2 

383 res = 32 

384 converged = 50 

385 

386 

387# --- Ra 1e6 --- 

388class RBC3DG4R4SDC44Ra1e6(RBC3DM4K4): 

389 Tend = 75 

390 dt = 1e-2 

391 res = 64 

392 converged = 22 

393 ic_config = {'config': RBC3DG4R4SDC34Ra1e5, 'res': 32, 'dt': 0.02} 

394 

395 

396class RBC3DG4R4SDC23Ra1e6(RBC3DM2K3): 

397 Tend = 75 

398 dt = 1e-2 

399 res = 64 

400 converged = 22 

401 ic_config = {'config': RBC3DG4R4SDC34Ra1e5, 'res': 32, 'dt': 0.02} 

402 

403 

404class RBC3DG4R4RKRa1e6(RBC3DverificationRK): 

405 Tend = 75 

406 dt = 1e-2 

407 res = 64 

408 ic_config = {'config': RBC3DG4R4SDC34Ra1e5, 'res': 32, 'dt': 0.02} 

409 converged = 22 

410 

411 

412class RBC3DG4R4EulerRa1e6(RBC3DverificationEuler): 

413 Tend = 75 

414 dt = 5e-3 

415 res = 64 

416 ic_config = {'config': RBC3DG4R4SDC34Ra1e5, 'res': 32, 'dt': 0.02} 

417 converged = 22 

418 

419 

420# --- Ra 1e7 --- 

421class RBC3DG4R4SDC23Ra1e7(RBC3DM2K3): 

422 Tend = 45 

423 dt = 5e-3 

424 res = 128 

425 converged = 25 

426 ic_config = {'config': RBC3DG4R4SDC23Ra1e6, 'res': 64, 'dt': 0.01} 

427 

428 

429class RBC3DG4R4SDC44Ra1e7(RBC3DM4K4): 

430 Tend = 45 

431 dt = 5e-3 

432 res = 128 

433 converged = 25 

434 ic_config = {'config': RBC3DG4R4SDC23Ra1e6, 'res': 64, 'dt': 0.01} 

435 

436 

437class RBC3DG4R4EulerRa1e7(RBC3DverificationEuler): 

438 Tend = 45 

439 dt = 1e-3 

440 res = 128 

441 converged = 25 

442 ic_config = {'config': RBC3DG4R4SDC23Ra1e6, 'res': 64, 'dt': 0.01} 

443 

444 

445class RBC3DG4R4RKRa1e7(RBC3DverificationRK): 

446 Tend = 45 

447 dt = 4e-3 

448 res = 128 

449 converged = 25 

450 ic_config = {'config': RBC3DG4R4SDC23Ra1e6, 'res': 64, 'dt': 0.01}