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

331 statements  

« prev     ^ index     » next       coverage.py v7.17.3, created at 2026-11-13 02:59 +0000

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

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 from pySDC.implementations.hooks.log_solution import LogToFile 

36 

37 return self.get_hook(LogToFile, filename=self.get_file_name(), time_increment=5e-1) 

38 

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

40 from pySDC.implementations.hooks.log_step_size import LogStepSize 

41 

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

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

44 return controller_params 

45 

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

47 from pySDC.implementations.problem_classes.RayleighBenard3D import ( 

48 RayleighBenard3D, 

49 ) 

50 from pySDC.implementations.problem_classes.generic_spectral import ( 

51 compute_residual_DAE, 

52 compute_residual_DAE_MPI, 

53 ) 

54 from pySDC.implementations.convergence_controller_classes.step_size_limiter import StepSizeSlopeLimiter 

55 from pySDC.implementations.convergence_controller_classes.crash import StopAtNan 

56 

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

58 

59 # subclass rather than patch: the sweeper class is shared with every other problem in the process 

60 sweeper_class = desc['sweeper_class'] 

61 desc['sweeper_class'] = type( 

62 sweeper_class.__name__, 

63 (sweeper_class,), 

64 {'compute_residual': compute_residual_DAE_MPI if MPIsweeper else compute_residual_DAE}, 

65 ) 

66 

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

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

69 

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

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

72 try: 

73 desc['convergence_controllers'].pop(LogStats) 

74 except KeyError: 

75 pass 

76 desc['convergence_controllers'][ClearStats] = {} 

77 

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

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

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

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

82 

83 res = 64 if res == -1 else res 

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

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

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

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

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

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

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

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

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

93 

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

95 

96 desc['problem_class'] = RayleighBenard3D 

97 

98 return desc 

99 

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

101 

102 if restart_idx == 0: 

103 P.logger.debug('Setting up random initial conditions') 

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

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

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

107 return u0_with_pressure, 0 

108 else: 

109 from pySDC.helpers.fieldsIO import FieldsIO 

110 

111 filename = self.get_file_name() 

112 P.logger.debug(f'Loading snapshot {restart_idx} from file {filename}') 

113 

114 P.setUpFieldsIO() 

115 outfile = FieldsIO.fromFile(filename) 

116 

117 t0, solution = outfile.readField(restart_idx) 

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

119 

120 if P.useGPU: 

121 solution = P.xp.array(solution) 

122 

123 u0 = P.u_init 

124 

125 if P.spectral_space: 

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

127 else: 

128 u0[...] = solution 

129 

130 P.logger.info(f'Loaded snapshot {restart_idx} at t={t0} from file {filename}') 

131 return u0, t0 

132 

133 def prepare_caches(self, prob): 

134 """ 

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

136 """ 

137 prob.logger.debug('Preparing caches ...') 

138 prob.eval_f(prob.u_init) 

139 prob.logger.debug('Prepared caches') 

140 

141 def prepare_for_benchmark(self): 

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

143 pass 

144 

145 self.get_LogToFile = _pass 

146 

147 def prepare_caches_for_benchmark(self, prob, controller): 

148 _rhs = prob.u_init 

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

150 _dt = sweeper.level.dt 

151 

152 hooks = controller.hooks 

153 

154 # mute controller 

155 type(controller).hooks = [] 

156 controller.run(_rhs, 0, _dt) 

157 

158 # unmute controller 

159 controller.hooks = hooks 

160 

161 try: 

162 import cuda as cp 

163 

164 cp.cuda.get_current_stream().synchronize() 

165 except ModuleNotFoundError: 

166 pass 

167 from mpi4py import MPI 

168 

169 MPI.COMM_WORLD.Barrier() 

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

171 

172 def prepare_description_for_benchmark(self, description, controller_params): 

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

174 from mpi4py import MPI 

175 

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

177 

178 controller_params['logger_level'] = 40 

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

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

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

182 

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

184 

185 # time_rank = 0 

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

187 # time_rank = description['sweeper_params']['comm'].rank 

188 # for i in range(MPI.COMM_WORLD.size): 

189 # if MPI.COMM_WORLD.rank == i: 

190 # print( 

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

192 # flush=True, 

193 # ) 

194 MPI.COMM_WORLD.barrier() 

195 

196 

197class RBC3Dverification(RayleighBenard3DRegular): 

198 converged = 0 

199 dt = 1e-2 

200 ic_config = { 

201 'config': None, 

202 'res': -1, 

203 'dt': -1, 

204 } 

205 res = None 

206 Ra = None 

207 Tend = 100 

208 res_ratio = 4 

209 gamma = 4 

210 

211 def get_file_name(self): 

212 res = self.args['res'] 

213 dt = self.args['dt'] 

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

215 

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

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

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

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

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

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

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

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

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

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

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

227 

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

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

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

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

232 

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

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

235 

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

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

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

239 return desc 

240 

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

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

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

244 

245 # read initial conditions 

246 from pySDC.helpers.fieldsIO import FieldsIO 

247 

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

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

250 ) 

251 ic_config.base_path = self.base_path 

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

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

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

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

256 

257 P.logger.debug( 

258 f'Setting up auxiliary problem with resolution {ic_nx}x{ic_ny}x{ic_nz} for interpolating initial conditions' 

259 ) 

260 _P = type(P)( 

261 nx=ic_nx, 

262 ny=ic_ny, 

263 nz=ic_nz, 

264 comm=P.comm, 

265 useGPU=P.useGPU, 

266 Dirichlet_recombination=False, 

267 left_preconditioner=False, 

268 ) 

269 _P.setUpFieldsIO() 

270 filename = ic_config.get_file_name() 

271 ic_file = FieldsIO.fromFile(filename) 

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

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

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

275 

276 # interpolate the initial conditions using padded transforms 

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

278 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}') 

279 

280 ics = _P.xp.array(ics) 

281 _ics_hat = _P.transform(ics) 

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

283 

284 self.get_LogToFile() 

285 

286 P.setUpFieldsIO() 

287 if P.spectral_space: 

288 u0_hat = P.u_init_forward 

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

290 return u0_hat, 0 

291 else: 

292 return ics_interpolated, 0 

293 

294 

295class RBC3DM2K3(RBC3Dverification): 

296 

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

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

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

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

301 return desc 

302 

303 

304class RBC3DM2K2(RBC3Dverification): 

305 

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

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

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

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

310 return desc 

311 

312 

313class RBC3DM3K4(RBC3Dverification): 

314 

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

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

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

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

319 return desc 

320 

321 

322class RBC3DM4K4(RBC3Dverification): 

323 

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

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

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

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

328 return desc 

329 

330 

331class RBC3DverificationRK(RBC3Dverification): 

332 

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

334 from pySDC.implementations.sweeper_classes.Runge_Kutta import ARK3 

335 

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

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

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

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

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

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

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

343 desc['sweeper_class'] = ARK3 

344 return desc 

345 

346 

347class RBC3DverificationEuler(RBC3DverificationRK): 

348 

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

350 from pySDC.implementations.sweeper_classes.Runge_Kutta import IMEXEulerStifflyAccurate 

351 

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

353 desc['sweeper_class'] = IMEXEulerStifflyAccurate 

354 return desc 

355 

356 

357class ClearStats(LogStats): 

358 

359 def post_step_processing(self, controller, S, **kwargs): 

360 self.reset_stats(controller) 

361 

362 def post_run_processing(self, controller, S, **kwargs): 

363 

364 def return_stats(): 

365 return {} 

366 

367 controller.return_stats = return_stats 

368 

369 

370# --- Ra 1e5 --- 

371class RBC3DG4R4SDC22Ra1e5(RBC3DM2K2): 

372 Tend = 200 

373 dt = 4e-2 

374 res = 32 

375 converged = 50 

376 

377 

378class RBC3DG4R4SDC23Ra1e5(RBC3DM2K3): 

379 Tend = 200 

380 dt = 6e-2 

381 res = 32 

382 converged = 50 

383 

384 

385class RBC3DG4R4SDC34Ra1e5(RBC3DM3K4): 

386 Tend = 200 

387 dt = 6e-2 

388 res = 32 

389 converged = 50 

390 

391 

392class RBC3DG4R4SDC44Ra1e5(RBC3DM4K4): 

393 Tend = 200 

394 dt = 6e-2 

395 res = 32 

396 converged = 50 

397 

398 

399class RBC3DG4R4RKRa1e5(RBC3DverificationRK): 

400 Tend = 200 

401 dt = 5e-2 

402 res = 32 

403 converged = 50 

404 

405 

406class RBC3DG4R4EulerRa1e5(RBC3DverificationEuler): 

407 Tend = 200 

408 dt = 2e-2 

409 res = 32 

410 converged = 50 

411 

412 

413# --- Ra 1e6 --- 

414class RBC3DG4R4SDC44Ra1e6(RBC3DM4K4): 

415 Tend = 75 

416 dt = 1e-2 

417 res = 64 

418 converged = 22 

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

420 

421 

422class RBC3DG4R4SDC23Ra1e6(RBC3DM2K3): 

423 Tend = 75 

424 dt = 1e-2 

425 res = 64 

426 converged = 22 

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

428 

429 

430class RBC3DG4R4RKRa1e6(RBC3DverificationRK): 

431 Tend = 75 

432 dt = 1e-2 

433 res = 64 

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

435 converged = 22 

436 

437 

438class RBC3DG4R4EulerRa1e6(RBC3DverificationEuler): 

439 Tend = 75 

440 dt = 5e-3 

441 res = 64 

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

443 converged = 22 

444 

445 

446# --- Ra 1e7 --- 

447class RBC3DG4R4SDC23Ra1e7(RBC3DM2K3): 

448 Tend = 45 

449 dt = 5e-3 

450 res = 128 

451 converged = 25 

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

453 

454 

455class RBC3DG4R4SDC44Ra1e7(RBC3DM4K4): 

456 Tend = 45 

457 dt = 5e-3 

458 res = 128 

459 converged = 25 

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

461 

462 

463class RBC3DG4R4EulerRa1e7(RBC3DverificationEuler): 

464 Tend = 45 

465 dt = 1e-3 

466 res = 128 

467 converged = 25 

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

469 

470 

471class RBC3DG4R4RKRa1e7(RBC3DverificationRK): 

472 Tend = 45 

473 dt = 4e-3 

474 res = 128 

475 converged = 25 

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

477 

478 

479# --- Ra 1e8 --- 

480class RBC3DG4R4SDC23Ra1e8(RBC3DM2K3): 

481 Tend = 70 

482 dt = 7e-4 

483 res = 256 

484 converged = 16 

485 ic_config = {'config': RBC3DG4R4SDC23Ra1e7, 'res': 128, 'dt': 0.005} 

486 

487 

488class RBC3DG4R4SDC44Ra1e8(RBC3DM4K4): 

489 Tend = 30 

490 dt = 5e-3 

491 res = 256 

492 converged = 8 

493 ic_config = {'config': RBC3DG4R4SDC23Ra1e7, 'res': 128, 'dt': 0.005} 

494 

495 

496class RBC3DG4R4EulerRa1e8(RBC3DverificationEuler): 

497 Tend = 30 

498 dt = 1e-4 

499 res = 256 

500 converged = 8 

501 ic_config = {'config': RBC3DG4R4SDC23Ra1e7, 'res': 128, 'dt': 0.005} 

502 

503 

504class RBC3DG4R4RKRa1e8(RBC3DverificationRK): 

505 Tend = 30 

506 dt = 4e-3 

507 res = 256 

508 converged = 8 

509 ic_config = {'config': RBC3DG4R4SDC23Ra1e7, 'res': 128, 'dt': 0.005}