Coverage for pySDC/implementations/problem_classes/RayleighBenard.py: 95%

257 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-09-29 12:50 +0000

1import numpy as np 

2from mpi4py import MPI 

3 

4from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear 

5from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh 

6from pySDC.core.convergence_controller import ConvergenceController 

7from pySDC.core.hooks import Hooks 

8from pySDC.core.check_convergence import CheckConvergence 

9from pySDC.core.problem import WorkCounter 

10 

11 

12class RayleighBenard(GenericSpectralLinear): 

13 """ 

14 2D Rayleigh-Benard convection, FFT in x and ultraspherical in z, IMEX with the nonlinear advection explicit. 

15 

16 Rayleigh-Benard Convection is a variation of incompressible Navier-Stokes. 

17 

18 The equations we solve are 

19 

20 u_x + v_z = 0 

21 T_t - kappa (T_xx + T_zz) = -uT_x - vT_z 

22 u_t - nu (u_xx + u_zz) + p_x = -uu_x - vu_z 

23 v_t - nu (v_xx + v_zz) + p_z - T = -uv_x - vv_z 

24 

25 with u the horizontal velocity, v the vertical velocity (in z-direction), T the temperature, p the pressure, indices 

26 denoting derivatives, kappa=(Rayleigh * Prandtl)**(-1/2) and nu = (Rayleigh / Prandtl)**(-1/2). Everything on the left 

27 hand side, that is the viscous part, the pressure gradient and the buoyancy due to temperature are treated 

28 implicitly, while the non-linear convection part on the right hand side is integrated explicitly. 

29 

30 The domain, vertical boundary conditions and pressure gauge are 

31 

32 Omega = [0, 8) x (-1, 1) 

33 T(z=+1) = 0 

34 T(z=-1) = 2 

35 u(z=+-1) = v(z=+-1) = 0 

36 integral over p = 0 

37 

38 The spectral discretization uses FFT horizontally, implying periodic BCs, and an ultraspherical method vertically to 

39 facilitate the Dirichlet BCs. 

40 

41 Parameters: 

42 Prandtl (float): Prandtl number 

43 Rayleigh (float): Rayleigh number 

44 nx (int): Horizontal resolution 

45 nz (int): Vertical resolution 

46 BCs (dict): Can specify boundary conditions here 

47 dealiasing (float): Dealiasing factor for evaluating the non-linear part 

48 comm (mpi4py.Intracomm): Space communicator 

49 """ 

50 

51 dtype_u = mesh 

52 dtype_f = imex_mesh 

53 

54 def __init__( 

55 self, 

56 Prandtl=1, 

57 Rayleigh=2e6, 

58 nx=256, 

59 nz=64, 

60 BCs=None, 

61 dealiasing=3 / 2, 

62 comm=None, 

63 Lx=4, 

64 Lz=1, 

65 z0=0, 

66 **kwargs, 

67 ): 

68 """ 

69 Constructor. `kwargs` are forwarded to parent class constructor. 

70 

71 Args: 

72 Prandtl (float): Prandtl number 

73 Rayleigh (float): Rayleigh number 

74 nx (int): Resolution in x-direction 

75 nz (int): Resolution in z direction 

76 BCs (dict): Vertical boundary conditions 

77 dealiasing (float): Dealiasing for evaluating the non-linear part in real space 

78 comm (mpi4py.Intracomm): Space communicator 

79 Lx (float): Horizontal length of the domain 

80 Lz (float): Vertical length of the domain 

81 z0 (float): Position of lower boundary 

82 """ 

83 BCs = {} if BCs is None else BCs 

84 BCs = { 

85 'T_top': 0, 

86 'T_bottom': 1, 

87 'v_top': 0, 

88 'v_bottom': 0, 

89 'u_top': 0, 

90 'u_bottom': 0, 

91 'p_integral': 0, 

92 **BCs, 

93 } 

94 if comm is None: 

95 try: 

96 from mpi4py import MPI 

97 

98 comm = MPI.COMM_WORLD 

99 except ModuleNotFoundError: 

100 pass 

101 self._makeAttributeAndRegister( 

102 'Prandtl', 

103 'Rayleigh', 

104 'nx', 

105 'nz', 

106 'BCs', 

107 'dealiasing', 

108 'comm', 

109 'Lx', 

110 'Lz', 

111 'z0', 

112 localVars=locals(), 

113 readOnly=True, 

114 ) 

115 

116 bases = [ 

117 {'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx}, 

118 {'base': 'ultraspherical', 'N': nz, 'x0': self.z0, 'x1': self.z0 + self.Lz}, 

119 ] 

120 components = ['u', 'v', 'T', 'p'] 

121 super().__init__(bases, components, comm=comm, **kwargs) 

122 

123 self.X, self.Z = self.get_grid() 

124 self.Kx, self.Kz = self.get_wavenumbers() 

125 

126 # construct 2D matrices 

127 Dzz = self.get_differentiation_matrix(axes=(1,), p=2) 

128 Dz = self.get_differentiation_matrix(axes=(1,)) 

129 Dx = self.get_differentiation_matrix(axes=(0,)) 

130 Dxx = self.get_differentiation_matrix(axes=(0,), p=2) 

131 Id = self.get_Id() 

132 

133 S1 = self.get_basis_change_matrix(p_out=0, p_in=1) 

134 S2 = self.get_basis_change_matrix(p_out=0, p_in=2) 

135 

136 U01 = self.get_basis_change_matrix(p_in=0, p_out=1) 

137 U12 = self.get_basis_change_matrix(p_in=1, p_out=2) 

138 U02 = self.get_basis_change_matrix(p_in=0, p_out=2) 

139 

140 self.Dx = Dx 

141 self.Dxx = Dxx 

142 self.Dz = S1 @ Dz 

143 self.Dzz = S2 @ Dzz 

144 

145 # compute rescaled Rayleigh number to extract viscosity and thermal diffusivity 

146 Ra = Rayleigh / (max([abs(BCs['T_top'] - BCs['T_bottom']), np.finfo(float).eps]) * self.axes[1].L ** 3) 

147 self.kappa = (Ra * Prandtl) ** (-1 / 2.0) 

148 self.nu = (Ra / Prandtl) ** (-1 / 2.0) 

149 

150 # construct operators 

151 L_lhs = { 

152 'p': {'u': U01 @ Dx, 'v': Dz}, # divergence free constraint 

153 'u': {'p': U02 @ Dx, 'u': -self.nu * (U02 @ Dxx + Dzz)}, 

154 'v': {'p': U12 @ Dz, 'v': -self.nu * (U02 @ Dxx + Dzz), 'T': -U02 @ Id}, 

155 'T': {'T': -self.kappa * (U02 @ Dxx + Dzz)}, 

156 } 

157 self.setup_L(L_lhs) 

158 

159 # mass matrix 

160 M_lhs = {i: {i: U02 @ Id} for i in ['u', 'v', 'T']} 

161 self.setup_M(M_lhs) 

162 

163 # Prepare going from second (first for divergence free equation) derivative basis back to Chebychov-T 

164 self.base_change = self._setup_operator({**{comp: {comp: S2} for comp in ['u', 'v', 'T']}, 'p': {'p': S1}}) 

165 

166 # BCs 

167 self.add_BC( 

168 component='p', equation='p', axis=1, v=self.BCs['p_integral'], kind='integral', line=-1, scalar=True 

169 ) 

170 self.add_BC(component='T', equation='T', axis=1, x=-1, v=self.BCs['T_bottom'], kind='Dirichlet', line=-1) 

171 self.add_BC(component='T', equation='T', axis=1, x=1, v=self.BCs['T_top'], kind='Dirichlet', line=-2) 

172 self.add_BC(component='v', equation='v', axis=1, x=1, v=self.BCs['v_top'], kind='Dirichlet', line=-1) 

173 self.add_BC(component='v', equation='v', axis=1, x=-1, v=self.BCs['v_bottom'], kind='Dirichlet', line=-2) 

174 self.remove_BC(component='v', equation='v', axis=1, x=-1, kind='Dirichlet', line=-2, scalar=True) 

175 self.add_BC(component='u', equation='u', axis=1, v=self.BCs['u_top'], x=1, kind='Dirichlet', line=-2) 

176 self.add_BC( 

177 component='u', 

178 equation='u', 

179 axis=1, 

180 v=self.BCs['u_bottom'], 

181 x=-1, 

182 kind='Dirichlet', 

183 line=-1, 

184 ) 

185 

186 # eliminate Nyquist mode if needed 

187 if nx % 2 == 0: 

188 Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index() 

189 for component in self.components: 

190 self.add_BC( 

191 component=component, equation=component, axis=0, kind='Nyquist', line=int(Nyquist_mode_index), v=0 

192 ) 

193 self.setup_BCs() 

194 

195 self.work_counters['rhs'] = WorkCounter() 

196 

197 def eval_f(self, u, *args, **kwargs): 

198 """ 

199 Evaluate the right hand side, split into an implicit and an explicit part. 

200 

201 The implicit part is -L u, i.e. diffusion, pressure gradient, buoyancy and, in the pressure line, the negative 

202 divergence, converted back to the Chebychev-T basis. The explicit part is the advection -(u d/dx + v d/dz) of u, 

203 v and T, computed in physical space on a grid padded by the dealiasing factor. 

204 

205 Args: 

206 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space 

207 *args: Not used, the right hand side does not depend on time 

208 **kwargs: Not used, the right hand side does not depend on time 

209 

210 Returns: 

211 dtype_f: The right hand side, in the same space as `u` 

212 """ 

213 f = self.f_init 

214 

215 if self.spectral_space: 

216 u_hat = u.copy() 

217 else: 

218 u_hat = self.transform(u) 

219 

220 f_impl_hat = self.u_init_forward 

221 

222 Dz = self.Dz 

223 Dx = self.Dx 

224 

225 iu, iv, iT, ip = self.index(['u', 'v', 'T', 'p']) 

226 

227 # evaluate implicit terms 

228 if not hasattr(self, '_L_T_base'): 

229 self._L_T_base = self.base_change @ self.L 

230 f_impl_hat = -(self._L_T_base @ u_hat.flatten()).reshape(u_hat.shape) 

231 

232 if self.spectral_space: 

233 f.impl[:] = f_impl_hat 

234 else: 

235 f.impl[:] = self.itransform(f_impl_hat).real 

236 

237 # ------------------------------------------- 

238 # treat convection explicitly with dealiasing 

239 

240 # start by computing derivatives 

241 if not hasattr(self, '_Dx_expanded') or not hasattr(self, '_Dz_expanded'): 

242 self._Dx_expanded = self._setup_operator({'u': {'u': Dx}, 'v': {'v': Dx}, 'T': {'T': Dx}, 'p': {}}) 

243 self._Dz_expanded = self._setup_operator({'u': {'u': Dz}, 'v': {'v': Dz}, 'T': {'T': Dz}, 'p': {}}) 

244 Dx_u_hat = (self._Dx_expanded @ u_hat.flatten()).reshape(u_hat.shape) 

245 Dz_u_hat = (self._Dz_expanded @ u_hat.flatten()).reshape(u_hat.shape) 

246 

247 padding = (self.dealiasing, self.dealiasing) 

248 Dx_u_pad = self.itransform(Dx_u_hat, padding=padding).real 

249 Dz_u_pad = self.itransform(Dz_u_hat, padding=padding).real 

250 u_pad = self.itransform(u_hat, padding=padding).real 

251 

252 fexpl_pad = self.xp.zeros_like(u_pad) 

253 fexpl_pad[iu][:] = -(u_pad[iu] * Dx_u_pad[iu] + u_pad[iv] * Dz_u_pad[iu]) 

254 fexpl_pad[iv][:] = -(u_pad[iu] * Dx_u_pad[iv] + u_pad[iv] * Dz_u_pad[iv]) 

255 fexpl_pad[iT][:] = -(u_pad[iu] * Dx_u_pad[iT] + u_pad[iv] * Dz_u_pad[iT]) 

256 

257 if self.spectral_space: 

258 f.expl[:] = self.transform(fexpl_pad, padding=padding) 

259 else: 

260 f.expl[:] = self.itransform(self.transform(fexpl_pad, padding=padding)).real 

261 

262 self.work_counters['rhs']() 

263 return f 

264 

265 def u_exact(self, t=0, noise_level=1e-3, seed=99): 

266 """ 

267 Initial conditions, which are only available at t=0. Velocities and temperature are linear in z between their 

268 boundary values, the pressure is zero, and the temperature is perturbed with seeded uniformly distributed noise, 

269 multiplied by `noise_level` and (z - z0) (z - z0 - Lz), which vanishes at both plates. 

270 

271 Args: 

272 t (float): Time, has to be 0 

273 noise_level (float): Amplitude of the noise 

274 seed (int): Seed for the random number generator 

275 

276 Returns: 

277 dtype_u: Initial conditions, in spectral space if `spectral_space` is set, else in physical space 

278 """ 

279 assert t == 0 

280 assert ( 

281 self.BCs['v_top'] == self.BCs['v_bottom'] 

282 ), 'Initial conditions are only implemented for zero velocity gradient' 

283 

284 me = self.spectral.u_init 

285 iu, iv, iT, ip = self.index(['u', 'v', 'T', 'p']) 

286 

287 # linear temperature gradient 

288 for comp in ['T', 'v', 'u']: 

289 a = (self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']) / self.Lz 

290 b = self.BCs[f'{comp}_bottom'] - a * self.z0 

291 me[self.index(comp)] = a * self.Z + b 

292 

293 # perturb slightly 

294 rng = self.xp.random.default_rng(seed=seed) 

295 

296 noise = self.spectral.u_init 

297 noise[iT] = rng.random(size=me[iT].shape) 

298 

299 me[iT] += noise[iT].real * noise_level * (self.Z - self.z0) * (self.Z - self.z0 - self.Lz) 

300 

301 if self.spectral_space: 

302 me_hat = self.spectral.u_init_forward 

303 me_hat[:] = self.transform(me) 

304 return me_hat 

305 else: 

306 return me 

307 

308 def apply_BCs(self, sol): 

309 """ 

310 Enforce the Dirichlet BCs at the top and bottom for arbitrary solution. 

311 The function modifies the last two modes of u, v, and T in order to achieve this. 

312 Note that the pressure is not modified here and the Nyquist mode is not altered either. 

313 

314 Args: 

315 sol: Some solution that does not need to enforce boundary conditions 

316 

317 Returns: 

318 Modified version of the solution that satisfies Dirichlet BCs. 

319 """ 

320 ultraspherical = self.spectral.axes[-1] 

321 

322 if self.spectral_space: 

323 sol_half_hat = self.itransform(sol, axes=(-2,)) 

324 else: 

325 sol_half_hat = self.transform(sol, axes=(-1,)) 

326 

327 BC_bottom = ultraspherical.get_BC(x=-1, kind='dirichlet') 

328 BC_top = ultraspherical.get_BC(x=1, kind='dirichlet') 

329 

330 M = np.array([BC_top[-2:], BC_bottom[-2:]]) 

331 M_I = np.linalg.inv(M) 

332 rhs = np.empty((2, self.nx), dtype=complex) 

333 for component in ['u', 'v', 'T']: 

334 i = self.index(component) 

335 rhs[0] = self.BCs[f'{component}_top'] - self.xp.sum(sol_half_hat[i, :, :-2] * BC_top[:-2], axis=1) 

336 rhs[1] = self.BCs[f'{component}_bottom'] - self.xp.sum(sol_half_hat[i, :, :-2] * BC_bottom[:-2], axis=1) 

337 

338 BC_vals = M_I @ rhs 

339 

340 sol_half_hat[i, :, -2:] = BC_vals.T 

341 

342 if self.spectral_space: 

343 return self.transform(sol_half_hat, axes=(-2,)) 

344 else: 

345 return self.itransform(sol_half_hat, axes=(-1,)) 

346 

347 def get_fig(self): # pragma: no cover 

348 """ 

349 Get a figure suitable to plot the solution of this problem 

350 

351 Returns 

352 ------- 

353 self.fig : matplotlib.pyplot.figure.Figure 

354 """ 

355 import matplotlib.pyplot as plt 

356 from mpl_toolkits.axes_grid1 import make_axes_locatable 

357 

358 self.fig, axs = plt.subplots(2, 1, sharex=True, sharey=True, figsize=((10, 5)), constrained_layout=True) 

359 self.cax = [] 

360 divider = make_axes_locatable(axs[0]) 

361 self.cax += [divider.append_axes('right', size='3%', pad=0.03)] 

362 divider2 = make_axes_locatable(axs[1]) 

363 self.cax += [divider2.append_axes('right', size='3%', pad=0.03)] 

364 return self.fig 

365 

366 def plot(self, u, t=None, fig=None, quantity='T'): # pragma: no cover 

367 r""" 

368 Plot the solution. 

369 

370 Parameters 

371 ---------- 

372 u : dtype_u 

373 Solution to be plotted 

374 t : float 

375 Time to display at the top of the figure 

376 fig : matplotlib.pyplot.figure.Figure 

377 Figure with the same structure as a figure generated by `self.get_fig`. If none is supplied, a new figure will be generated. 

378 quantity : (str) 

379 quantity you want to plot 

380 

381 Returns 

382 ------- 

383 None 

384 """ 

385 fig = self.get_fig() if fig is None else fig 

386 axs = fig.axes 

387 

388 imV = axs[1].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real) 

389 

390 if self.spectral_space: 

391 u = self.itransform(u) 

392 

393 imT = axs[0].pcolormesh(self.X, self.Z, u[self.index(quantity)].real) 

394 

395 for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity'], strict=True): 

396 axs[i].set_aspect(1) 

397 axs[i].set_title(label) 

398 

399 if t is not None: 

400 fig.suptitle(f't = {t:.2f}') 

401 axs[1].set_xlabel(r'$x$') 

402 axs[1].set_ylabel(r'$z$') 

403 fig.colorbar(imT, self.cax[0]) 

404 fig.colorbar(imV, self.cax[1]) 

405 

406 def compute_vorticity(self, u): 

407 """ 

408 Compute the vorticity by spectral differentiation, as Dx v + Dz u. 

409 

410 Args: 

411 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space 

412 

413 Returns: 

414 xp.ndarray: Vorticity in physical space 

415 """ 

416 if self.spectral_space: 

417 u_hat = u.copy() 

418 else: 

419 u_hat = self.transform(u) 

420 

421 Dz = self.Dz 

422 Dx = self.Dx 

423 iu, iv = self.index(['u', 'v']) 

424 

425 vorticity_hat = self.spectral.u_init_forward 

426 vorticity_hat[0] = (Dx * u_hat[iv].flatten() - Dz @ u_hat[iu].flatten()).reshape(u_hat[iu].shape) 

427 return self.itransform(vorticity_hat)[0].real 

428 

429 def getOutputFile(self, fileName): 

430 """ 

431 Set up a `Rectilinear` output file on the grid of this problem, with one variable per component plus the 

432 vorticity, see `processSolutionForOutput`. 

433 

434 Args: 

435 fileName (str): Name of the file 

436 

437 Returns: 

438 pySDC.helpers.fieldsIO.Rectilinear: The initialized output file 

439 """ 

440 from pySDC.helpers.fieldsIO import Rectilinear 

441 

442 self.setUpFieldsIO() 

443 

444 coords = [me.get_1dgrid() for me in self.spectral.axes] 

445 assert np.allclose([len(me) for me in coords], self.spectral.global_shape[1:]) 

446 

447 fOut = Rectilinear(np.float64, fileName=fileName) 

448 fOut.setHeader(nVar=len(self.components) + 1, coords=coords) 

449 fOut.initialize() 

450 return fOut 

451 

452 def processSolutionForOutput(self, u): 

453 """ 

454 Prepare a solution for output: the real parts of all components in physical space, with the vorticity appended 

455 as the last variable. 

456 

457 Args: 

458 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space 

459 

460 Returns: 

461 numpy.ndarray: Components u, v, T, p and the vorticity 

462 """ 

463 vorticity = self.compute_vorticity(u) 

464 

465 if self.spectral_space: 

466 u_real = self.itransform(u).real 

467 else: 

468 u_real = u.real 

469 

470 me = np.empty(shape=(u_real.shape[0] + 1, *vorticity.shape)) 

471 me[:-1] = u_real 

472 me[-1] = vorticity 

473 return me 

474 

475 def compute_Nusselt_numbers(self, u): 

476 """ 

477 Compute the various versions of the Nusselt number. This reflects the type of heat transport. 

478 If the Nusselt number is equal to one, it indicates heat transport due to conduction. If it is larger, 

479 advection is present. 

480 Computing the Nusselt number at various places can be used to check the code. 

481 

482 Args: 

483 u: The solution you want to compute the Nusselt numbers of 

484 

485 Returns: 

486 dict: Nusselt number averaged over the entire volume and horizontally averaged at the top and bottom. 

487 """ 

488 iv, iT = self.index(['v', 'T']) 

489 zAxis = self.spectral.axes[-1] 

490 

491 if self.spectral_space: 

492 u_hat = u.copy() 

493 else: 

494 u_hat = self.transform(u) 

495 

496 DzT_hat = (self.Dz @ u_hat[iT].flatten()).reshape(u_hat[iT].shape) 

497 

498 # compute vT with dealiasing 

499 padding = (self.dealiasing, self.dealiasing) 

500 u_pad = self.itransform(u_hat, padding=padding).real 

501 _me = self.xp.zeros_like(u_pad) 

502 _me[0] = u_pad[iv] * u_pad[iT] 

503 vT_hat = self.transform(_me, padding=padding)[0] 

504 

505 if not hasattr(self, '_zInt'): 

506 self._zInt = zAxis.get_integration_matrix() 

507 

508 nusselt_hat = (vT_hat / self.kappa - DzT_hat) * self.axes[-1].L 

509 

510 # get coefficients for evaluation on the boundary 

511 top = zAxis.get_BC(kind='Dirichlet', x=1) 

512 bot = zAxis.get_BC(kind='Dirichlet', x=-1) 

513 

514 integral_V = 0 

515 if self.comm.rank == 0: 

516 

517 integral_z = (self._zInt @ nusselt_hat[0]).real 

518 integral_z[0] = zAxis.get_integration_constant(integral_z, axis=-1) 

519 integral_V = ((top - bot) * integral_z).sum() * self.axes[0].L / self.nx 

520 

521 Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0) 

522 Nusselt_t = self.comm.bcast(self.xp.sum(nusselt_hat.real[0] * top, axis=-1) / self.nx, root=0) 

523 Nusselt_b = self.comm.bcast(self.xp.sum(nusselt_hat.real[0] * bot, axis=-1) / self.nx, root=0) 

524 

525 return { 

526 'V': Nusselt_V, 

527 't': Nusselt_t, 

528 'b': Nusselt_b, 

529 } 

530 

531 def compute_viscous_dissipation(self, u): 

532 """ 

533 Compute the pointwise viscous dissipation ``abs(u Lap(u) + v Lap(v))``, the velocity times its Laplacian, where 

534 the 

535 Laplacian is computed by spectral differentiation. 

536 

537 Args: 

538 u (dtype_u): Solution in physical space 

539 

540 Returns: 

541 xp.ndarray: Viscous dissipation in physical space 

542 """ 

543 iu, iv = self.index(['u', 'v']) 

544 

545 Lap_u_hat = self.spectral.u_init_forward 

546 

547 if self.spectral_space: 

548 u_hat = u.copy() 

549 u = self.spectral.u_init 

550 u[...] = self.itransform(u_hat).real 

551 else: 

552 u_hat = self.transform(u) 

553 Lap_u_hat[iu] = ((self.Dzz + self.Dxx) @ u_hat[iu].flatten()).reshape(u_hat[iu].shape) 

554 Lap_u_hat[iv] = ((self.Dzz + self.Dxx) @ u_hat[iv].flatten()).reshape(u_hat[iu].shape) 

555 Lap_u = self.itransform(Lap_u_hat) 

556 

557 return abs(u[iu] * Lap_u[iu] + u[iv] * Lap_u[iv]) 

558 

559 def compute_buoyancy_generation(self, u): 

560 """ 

561 Compute the pointwise buoyancy production ``abs(Rayleigh v T)``. 

562 

563 Args: 

564 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space 

565 

566 Returns: 

567 xp.ndarray: Buoyancy production in physical space 

568 """ 

569 if self.spectral_space: 

570 u = self.itransform(u) 

571 iv, iT = self.index(['v', 'T']) 

572 return abs(u[iv] * self.Rayleigh * u[iT]) 

573 

574 

575class CFLLimit(ConvergenceController): 

576 

577 def dependencies(self, controller, *args, **kwargs): 

578 """ 

579 Add the hooks `LogCFL` and `LogStepSize` to the controller to record the CFL limit and the step size. 

580 

581 Args: 

582 controller (pySDC.Controller): The controller 

583 *args: Not used, e.g. the description 

584 **kwargs: Not used, e.g. the description 

585 

586 Returns: 

587 None 

588 """ 

589 from pySDC.implementations.hooks.log_step_size import LogStepSize 

590 

591 controller.add_hook(LogCFL) 

592 controller.add_hook(LogStepSize) 

593 

594 def setup_status_variables(self, controller, **kwargs): 

595 """ 

596 Add the embedded error variable to the error function. 

597 

598 Args: 

599 controller (pySDC.Controller): The controller 

600 """ 

601 self.add_status_variable_to_level('CFL_limit') 

602 

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

604 """ 

605 Define default parameters here. 

606 

607 Default parameters are: 

608 - control_order (int): The order relative to other convergence controllers 

609 - dt_max (float): maximal step size 

610 - dt_min (float): minimal step size 

611 

612 Args: 

613 controller (pySDC.Controller): The controller 

614 params (dict): The params passed for this specific convergence controller 

615 description (dict): The description object used to instantiate the controller 

616 

617 Returns: 

618 (dict): The updated params dictionary 

619 """ 

620 defaults = { 

621 "control_order": -50, 

622 "dt_max": np.inf, 

623 "dt_min": 0, 

624 "cfl": 0.4, 

625 } 

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

627 

628 @staticmethod 

629 def compute_max_step_size(P, u): 

630 """ 

631 Compute the largest step size allowed by the CFL condition with CFL number 1: the minimum over the grid of the 

632 grid spacing divided by the absolute velocity, in x and in z, reduced over the space communicator of the 

633 problem. The vertical grid spacing is the distance between the midpoints of neighbouring Chebychev nodes, with 

634 the domain [z0, z0 + Lz]. 

635 

636 Args: 

637 P (RayleighBenard): The problem 

638 u (dtype_u): Solution, in spectral space if `P.spectral_space` is set, else in physical space 

639 

640 Returns: 

641 float: Maximal step size 

642 """ 

643 grid_spacing_x = P.X[1, 0] - P.X[0, 0] 

644 

645 cell_wallz = P.xp.zeros(P.nz + 1) 

646 cell_wallz[0] = P.z0 + P.Lz 

647 cell_wallz[-1] = P.z0 

648 cell_wallz[1:-1] = (P.Z[0, :-1] + P.Z[0, 1:]) / 2 

649 grid_spacing_z = cell_wallz[:-1] - cell_wallz[1:] 

650 

651 iu, iv = P.index(['u', 'v']) 

652 

653 if P.spectral_space: 

654 u = P.itransform(u) 

655 

656 max_step_size_x = P.xp.min(grid_spacing_x / P.xp.abs(u[iu])) 

657 max_step_size_z = P.xp.min(grid_spacing_z / P.xp.abs(u[iv])) 

658 max_step_size = min([max_step_size_x, max_step_size_z]) 

659 

660 if hasattr(P, 'comm'): 

661 max_step_size = P.comm.allreduce(max_step_size, op=MPI.MIN) 

662 return float(max_step_size) 

663 

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

665 """ 

666 Once the step has converged, compute the CFL limit at the end point of the finest level, store `cfl` times it in 

667 the level status for `LogCFL` and limit the new step size to it. If no other convergence controller has set a 

668 new step size, max(`dt_max`, current step size) is limited instead. The result is at least `dt_min`. 

669 

670 Args: 

671 controller (pySDC.Controller): The controller 

672 step (pySDC.Step): The current step 

673 

674 Returns: 

675 None 

676 """ 

677 if not CheckConvergence.check_convergence(step): 

678 return None 

679 

680 L = step.levels[0] 

681 P = step.levels[0].prob 

682 

683 L.sweep.compute_end_point() 

684 max_step_size = self.compute_max_step_size(P, L.uend) 

685 

686 L.status.CFL_limit = self.params.cfl * max_step_size 

687 

688 dt_new = L.status.dt_new if L.status.dt_new else self.params.dt_max 

689 L.status.dt_new = min([dt_new, self.params.dt_max, self.params.cfl * max_step_size]) 

690 L.status.dt_new = max([self.params.dt_min, L.status.dt_new]) 

691 

692 self.log(f'dt max: {max_step_size:.2e} -> New step size: {L.status.dt_new:.2e}', step) 

693 

694 

695class LogCFL(Hooks): 

696 

697 def post_step(self, step, level_number): 

698 """ 

699 Record CFL limit. 

700 

701 Args: 

702 step (pySDC.Step.step): the current step 

703 level_number (int): the current level number 

704 

705 Returns: 

706 None 

707 """ 

708 super().post_step(step, level_number) 

709 

710 L = step.levels[level_number] 

711 

712 self.add_to_stats( 

713 process=step.status.slot, 

714 time=L.time + L.dt, 

715 level=L.level_index, 

716 iter=step.status.iter, 

717 sweep=L.status.sweep, 

718 type='CFL_limit', 

719 value=L.status.CFL_limit, 

720 ) 

721 

722 

723class LogAnalysisVariables(Hooks): 

724 

725 def post_step(self, step, level_number): 

726 """ 

727 Record Nusselt numbers. 

728 

729 Args: 

730 step (pySDC.Step.step): the current step 

731 level_number (int): the current level number 

732 

733 Returns: 

734 None 

735 """ 

736 super().post_step(step, level_number) 

737 

738 L = step.levels[level_number] 

739 P = L.prob 

740 

741 L.sweep.compute_end_point() 

742 Nusselt = P.compute_Nusselt_numbers(L.uend) 

743 buoyancy_production = P.compute_buoyancy_generation(L.uend) 

744 viscous_dissipation = P.compute_viscous_dissipation(L.uend) 

745 

746 for key, value in zip( 

747 ['Nusselt', 'buoyancy_production', 'viscous_dissipation'], 

748 [Nusselt, buoyancy_production, viscous_dissipation], 

749 strict=True, 

750 ): 

751 self.add_to_stats( 

752 process=step.status.slot, 

753 time=L.time + L.dt, 

754 level=L.level_index, 

755 iter=step.status.iter, 

756 sweep=L.status.sweep, 

757 type=key, 

758 value=value, 

759 )