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

295 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-09-28 13:30 +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 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 

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

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

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

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

77 

78 res = 64 if res == -1 else res 

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

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

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

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

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

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

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

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

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

88 

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

90 

91 desc['problem_class'] = RayleighBenard3D 

92 

93 return desc 

94 

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

96 

97 if restart_idx == 0: 

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

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

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

101 return u0_with_pressure, 0 

102 else: 

103 from pySDC.helpers.fieldsIO import FieldsIO 

104 

105 P.setUpFieldsIO() 

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

107 

108 t0, solution = outfile.readField(restart_idx) 

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

110 

111 if P.useGPU: 

112 solution = P.xp.array(solution) 

113 

114 u0 = P.u_init 

115 

116 if P.spectral_space: 

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

118 else: 

119 u0[...] = solution 

120 

121 return u0, t0 

122 

123 def prepare_caches(self, prob): 

124 """ 

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

126 """ 

127 prob.eval_f(prob.u_init) 

128 

129 def prepare_for_benchmark(self): 

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

131 pass 

132 

133 self.get_LogToFile = _pass 

134 

135 def prepare_caches_for_benchmark(self, prob, controller): 

136 _rhs = prob.u_init 

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

138 _dt = sweeper.level.dt 

139 

140 hooks = controller.hooks 

141 

142 # mute controller 

143 type(controller).hooks = [] 

144 controller.run(_rhs, 0, _dt) 

145 

146 # unmute controller 

147 controller.hooks = hooks 

148 

149 try: 

150 import cuda as cp 

151 

152 cp.cuda.get_current_stream().synchronize() 

153 except ModuleNotFoundError: 

154 pass 

155 from mpi4py import MPI 

156 

157 MPI.COMM_WORLD.Barrier() 

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

159 

160 def prepare_description_for_benchmark(self, description, controller_params): 

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

162 from mpi4py import MPI 

163 

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

165 

166 controller_params['logger_level'] = 40 

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

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

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

170 

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

172 

173 time_rank = 0 

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

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

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

177 if MPI.COMM_WORLD.rank == i: 

178 print( 

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

180 flush=True, 

181 ) 

182 MPI.COMM_WORLD.barrier() 

183 

184 

185class RBC3Dverification(RayleighBenard3DRegular): 

186 converged = 0 

187 dt = 1e-2 

188 ic_config = { 

189 'config': None, 

190 'res': -1, 

191 'dt': -1, 

192 } 

193 res = None 

194 Ra = None 

195 Tend = 100 

196 res_ratio = 4 

197 gamma = 4 

198 

199 def get_file_name(self): 

200 res = self.args['res'] 

201 dt = self.args['dt'] 

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

203 

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

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

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

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

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

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

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

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

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

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

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

215 

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

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

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

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

220 

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

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

223 

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

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

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

227 return desc 

228 

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

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

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

232 

233 # read initial conditions 

234 from pySDC.helpers.fieldsIO import FieldsIO 

235 

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

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

238 ) 

239 ic_config.base_path = self.base_path 

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

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

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

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

244 

245 _P = type(P)( 

246 nx=ic_nx, 

247 ny=ic_ny, 

248 nz=ic_nz, 

249 comm=P.comm, 

250 useGPU=P.useGPU, 

251 Dirichlet_recombination=False, 

252 left_preconditioner=False, 

253 ) 

254 _P.setUpFieldsIO() 

255 filename = ic_config.get_file_name() 

256 ic_file = FieldsIO.fromFile(filename) 

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

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

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

260 

261 # interpolate the initial conditions using padded transforms 

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

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

264 

265 ics = _P.xp.array(ics) 

266 _ics_hat = _P.transform(ics) 

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

268 

269 self.get_LogToFile() 

270 

271 P.setUpFieldsIO() 

272 if P.spectral_space: 

273 u0_hat = P.u_init_forward 

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

275 return u0_hat, 0 

276 else: 

277 return ics_interpolated, 0 

278 

279 

280class RBC3DM2K3(RBC3Dverification): 

281 

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

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

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

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

286 return desc 

287 

288 

289class RBC3DM2K2(RBC3Dverification): 

290 

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

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

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

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

295 return desc 

296 

297 

298class RBC3DM3K4(RBC3Dverification): 

299 

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

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

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

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

304 return desc 

305 

306 

307class RBC3DM4K4(RBC3Dverification): 

308 

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

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

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

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

313 return desc 

314 

315 

316class RBC3DverificationRK(RBC3Dverification): 

317 

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

319 from pySDC.implementations.sweeper_classes.Runge_Kutta import ARK3 

320 

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

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

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

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

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

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

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

328 desc['sweeper_class'] = ARK3 

329 return desc 

330 

331 

332class RBC3DverificationEuler(RBC3DverificationRK): 

333 

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

335 from pySDC.implementations.sweeper_classes.Runge_Kutta import IMEXEulerStifflyAccurate 

336 

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

338 desc['sweeper_class'] = IMEXEulerStifflyAccurate 

339 return desc 

340 

341 

342# --- Ra 1e5 --- 

343class RBC3DG4R4SDC22Ra1e5(RBC3DM2K2): 

344 Tend = 200 

345 dt = 4e-2 

346 res = 32 

347 converged = 50 

348 

349 

350class RBC3DG4R4SDC23Ra1e5(RBC3DM2K3): 

351 Tend = 200 

352 dt = 6e-2 

353 res = 32 

354 converged = 50 

355 

356 

357class RBC3DG4R4SDC34Ra1e5(RBC3DM3K4): 

358 Tend = 200 

359 dt = 6e-2 

360 res = 32 

361 converged = 50 

362 

363 

364class RBC3DG4R4SDC44Ra1e5(RBC3DM4K4): 

365 Tend = 200 

366 dt = 6e-2 

367 res = 32 

368 converged = 50 

369 

370 

371class RBC3DG4R4RKRa1e5(RBC3DverificationRK): 

372 Tend = 200 

373 dt = 5e-2 

374 res = 32 

375 converged = 50 

376 

377 

378class RBC3DG4R4EulerRa1e5(RBC3DverificationEuler): 

379 Tend = 200 

380 dt = 2e-2 

381 res = 32 

382 converged = 50 

383 

384 

385# --- Ra 1e6 --- 

386class RBC3DG4R4SDC44Ra1e6(RBC3DM4K4): 

387 Tend = 75 

388 dt = 1e-2 

389 res = 64 

390 converged = 22 

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

392 

393 

394class RBC3DG4R4SDC23Ra1e6(RBC3DM2K3): 

395 Tend = 75 

396 dt = 1e-2 

397 res = 64 

398 converged = 22 

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

400 

401 

402class RBC3DG4R4RKRa1e6(RBC3DverificationRK): 

403 Tend = 75 

404 dt = 1e-2 

405 res = 64 

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

407 converged = 22 

408 

409 

410class RBC3DG4R4EulerRa1e6(RBC3DverificationEuler): 

411 Tend = 75 

412 dt = 5e-3 

413 res = 64 

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

415 converged = 22 

416 

417 

418# --- Ra 1e7 --- 

419class RBC3DG4R4SDC23Ra1e7(RBC3DM2K3): 

420 Tend = 45 

421 dt = 5e-3 

422 res = 128 

423 converged = 25 

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

425 

426 

427class RBC3DG4R4SDC44Ra1e7(RBC3DM4K4): 

428 Tend = 45 

429 dt = 5e-3 

430 res = 128 

431 converged = 25 

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

433 

434 

435class RBC3DG4R4EulerRa1e7(RBC3DverificationEuler): 

436 Tend = 45 

437 dt = 1e-3 

438 res = 128 

439 converged = 25 

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

441 

442 

443class RBC3DG4R4RKRa1e7(RBC3DverificationRK): 

444 Tend = 45 

445 dt = 4e-3 

446 res = 128 

447 converged = 25 

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