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

1import numpy as np

2from mpi4py import MPI

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.implementations.convergence_controller_classes.check_convergence import CheckConvergence

9from pySDC.core.problem import WorkCounter

12class RayleighBenard(GenericSpectralLinear):

13 """

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

16 The equations we solve are

18 u_x + v_z = 0

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

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

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

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

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

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

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

28 The domain, vertical boundary conditions and pressure gauge are

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

31 T(z=+1) = 0

32 T(z=-1) = 2

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

34 integral over p = 0

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

37 facilitate the Dirichlet BCs.

39 Parameters:

40 Prandtl (float): Prandtl number

41 Rayleigh (float): Rayleigh number

42 nx (int): Horizontal resolution

43 nz (int): Vertical resolution

44 BCs (dict): Can specify boundary conditions here

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

46 comm (mpi4py.Intracomm): Space communicator

47 """

49 dtype_u = mesh

50 dtype_f = imex_mesh

52 def __init__(

53 self,

54 Prandtl=1,

55 Rayleigh=2e6,

56 nx=256,

57 nz=64,

58 BCs=None,

59 dealiasing=3 / 2,

60 comm=None,

61 Lx=8,

62 **kwargs,

63 ):

64 """

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

67 Args:

68 Prandtl (float): Prandtl number

69 Rayleigh (float): Rayleigh number

70 nx (int): Resolution in x-direction

71 nz (int): Resolution in z direction

72 BCs (dict): Vertical boundary conditions

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

74 comm (mpi4py.Intracomm): Space communicator

75 Lx (float): Horizontal length of the domain

76 """

77 BCs = {} if BCs is None else BCs

78 BCs = {

79 'T_top': 0,

80 'T_bottom': 2,

81 'v_top': 0,

82 'v_bottom': 0,

83 'u_top': 0,

84 'u_bottom': 0,

85 'p_integral': 0,

86 **BCs,

87 }

88 if comm is None:

89 try:

90 from mpi4py import MPI

92 comm = MPI.COMM_WORLD

93 except ModuleNotFoundError:

94 pass

95 self._makeAttributeAndRegister(

96 'Prandtl',

97 'Rayleigh',

98 'nx',

99 'nz',

100 'BCs',

101 'dealiasing',

102 'comm',

103 'Lx',

104 localVars=locals(),

105 readOnly=True,

106 )

107

108 bases = [{'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx}, {'base': 'ultraspherical', 'N': nz}]

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

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

111

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

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

114

115 # construct 2D matrices

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

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

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

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

120 Id = self.get_Id()

121

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

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

124

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

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

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

128

129 self.Dx = Dx

130 self.Dxx = Dxx

131 self.Dz = S1 @ Dz

132 self.Dzz = S2 @ Dzz

133

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

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

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

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

138

139 # construct operators

140 L_lhs = {

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

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

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

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

145 }

146 self.setup_L(L_lhs)

147

148 # mass matrix

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

150 self.setup_M(M_lhs)

151

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

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

154

155 # BCs

156 self.add_BC(

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

158 )

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

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

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

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

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

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

165 self.add_BC(

166 component='u',

167 equation='u',

168 axis=1,

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

170 x=-1,

171 kind='Dirichlet',

172 line=-1,

173 )

174

175 # eliminate Nyquist mode if needed

176 if nx % 2 == 0:

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

178 for component in self.components:

179 self.add_BC(

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

181 )

182 self.setup_BCs()

183

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

185

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

187 f = self.f_init

188

189 if self.spectral_space:

190 u_hat = u.copy()

191 else:

192 u_hat = self.transform(u)

193

194 f_impl_hat = self.u_init_forward

195

196 Dz = self.Dz

197 Dx = self.Dx

198

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

200

201 # evaluate implicit terms

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

203 self._L_T_base = self.base_change @ self.L

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

205

206 if self.spectral_space:

207 f.impl[:] = f_impl_hat

208 else:

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

210

211 # -------------------------------------------

212 # treat convection explicitly with dealiasing

213

214 # start by computing derivatives

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

216 self._Dx_expanded = self._setup_operator(

217 {'u': {'u': Dx}, 'v': {'v': Dx}, 'T': {'T': Dx}, 'p': {}}, diag=True

218 )

219 self._Dz_expanded = self._setup_operator(

220 {'u': {'u': Dz}, 'v': {'v': Dz}, 'T': {'T': Dz}, 'p': {}}, diag=True

221 )

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

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

224

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

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

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

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

229

230 fexpl_pad = self.xp.zeros_like(u_pad)

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

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

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

234

235 if self.spectral_space:

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

237 else:

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

239

240 self.work_counters['rhs']()

241 return f

242

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

244 assert t == 0

245 assert (

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

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

248

249 me = self.spectral.u_init

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

251

252 # linear temperature gradient

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

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

255 b = (self.BCs[f'{comp}_top'] + self.BCs[f'{comp}_bottom']) / 2

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

257

258 # perturb slightly

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

260

261 noise = self.spectral.u_init

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

263

264 me[iT] += noise[iT].real * noise_level * (self.Z - 1) * (self.Z + 1)

265

266 if self.spectral_space:

267 me_hat = self.spectral.u_init_forward

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

269 return me_hat

270 else:

271 return me

272

273 def apply_BCs(self, sol):

274 """

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

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

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

278

279 Args:

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

281

282 Returns:

283 Modified version of the solution that satisfies Dirichlet BCs.

284 """

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

286

287 if self.spectral_space:

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

289 else:

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

291

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

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

294

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

296 M_I = np.linalg.inv(M)

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

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

299 i = self.index(component)

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

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

302

303 BC_vals = M_I @ rhs

304

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

306

307 if self.spectral_space:

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

309 else:

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

311

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

313 """

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

315

316 Returns

317 -------

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

319 """

320 import matplotlib.pyplot as plt

321 from mpl_toolkits.axes_grid1 import make_axes_locatable

322

323 plt.rcParams['figure.constrained_layout.use'] = True

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

325 self.cax = []

326 divider = make_axes_locatable(axs[0])

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

328 divider2 = make_axes_locatable(axs[1])

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

330 return self.fig

331

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

333 r"""

334 Plot the solution.

335

336 Parameters

337 ----------

338 u : dtype_u

339 Solution to be plotted

340 t : float

341 Time to display at the top of the figure

342 fig : matplotlib.pyplot.figure.Figure

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

344 quantity : (str)

345 quantity you want to plot

346

347 Returns

348 -------

349 None

350 """

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

352 axs = fig.axes

353

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

355

356 if self.spectral_space:

357 u = self.itransform(u)

358

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

360

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

362 axs[i].set_aspect(1)

363 axs[i].set_title(label)

364

365 if t is not None:

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

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

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

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

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

371

372 def compute_vorticity(self, u):

373 if self.spectral_space:

374 u_hat = u.copy()

375 else:

376 u_hat = self.transform(u)

377 Dz = self.Dz

378 Dx = self.Dx

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

380

381 vorticity_hat = self.spectral.u_init_forward

382 vorticity_hat[0] = (Dx * u_hat[iv].flatten() + Dz @ u_hat[iu].flatten()).reshape(u[iu].shape)

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

384

385 def compute_Nusselt_numbers(self, u):

386 """

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

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

389 advection is present.

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

391

392 Args:

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

394

395 Returns:

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

397 """

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

399

400 DzT_hat = self.spectral.u_init_forward

401

402 if self.spectral_space:

403 u_hat = u.copy()

404 else:

405 u_hat = self.transform(u)

406

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

408

409 # compute vT with dealiasing

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

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

412 _me = self.xp.zeros_like(u_pad)

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

414 vT_hat = self.transform(_me, padding=padding)

415

416 nusselt_hat = (vT_hat[0] - DzT_hat[iT]) / self.nx

417 nusselt_no_v_hat = (-DzT_hat[iT]) / self.nx

418

419 integral_z = self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='integral'), axis=-1).real

420 integral_V = (

421 integral_z[0] * self.axes[0].L

422 ) # only the first Fourier mode has non-zero integral with periodic BCs

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

424

425 Nusselt_t = self.comm.bcast(

426 self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=1), axis=-1).real[0], root=0

427 )

428 Nusselt_b = self.comm.bcast(

429 self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=-1), axis=-1).real[0], root=0

430 )

431 Nusselt_no_v_t = self.comm.bcast(

432 self.xp.sum(nusselt_no_v_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=1), axis=-1).real[0], root=0

433 )

434 Nusselt_no_v_b = self.comm.bcast(

435 self.xp.sum(nusselt_no_v_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=-1), axis=-1).real[0],

436 root=0,

437 )

438

439 return {

440 'V': Nusselt_V,

441 't': Nusselt_t,

442 'b': Nusselt_b,

443 't_no_v': Nusselt_no_v_t,

444 'b_no_v': Nusselt_no_v_b,

445 }

446

447 def compute_viscous_dissipation(self, u):

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

449

450 Lap_u_hat = self.spectral.u_init_forward

451

452 if self.spectral_space:

453 u_hat = u.copy()

454 else:

455 u_hat = self.transform(u)

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

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

458 Lap_u = self.itransform(Lap_u_hat)

459

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

461

462 def compute_buoyancy_generation(self, u):

463 if self.spectral_space:

464 u = self.itransform(u)

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

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

467

468

469class CFLLimit(ConvergenceController):

470

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

472 from pySDC.implementations.hooks.log_step_size import LogStepSize

473

474 controller.add_hook(LogCFL)

475 controller.add_hook(LogStepSize)

476

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

478 """

479 Add the embedded error variable to the error function.

480

481 Args:

482 controller (pySDC.Controller): The controller

483 """

484 self.add_status_variable_to_level('CFL_limit')

485

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

487 """

488 Define default parameters here.

489

490 Default parameters are:

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

492 - dt_max (float): maximal step size

493 - dt_min (float): minimal step size

494

495 Args:

496 controller (pySDC.Controller): The controller

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

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

499

500 Returns:

501 (dict): The updated params dictionary

502 """

503 defaults = {

504 "control_order": -50,

505 "dt_max": np.inf,

506 "dt_min": 0,

507 "cfl": 0.4,

508 }

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

510

511 @staticmethod

512 def compute_max_step_size(P, u):

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

514

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

516 cell_wallz[0] = 1

517 cell_wallz[-1] = -1

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

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

520

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

522

523 if P.spectral_space:

524 u = P.itransform(u)

525

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

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

528 max_step_size = min([max_step_size_x, max_step_size_z])

529

530 if hasattr(P, 'comm'):

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

532 return float(max_step_size)

533

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

535 if not CheckConvergence.check_convergence(step):

536 return None

537

538 L = step.levels[0]

539 P = step.levels[0].prob

540

541 L.sweep.compute_end_point()

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

543

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

545

546 dt_new = L.status.dt_new if L.status.dt_new else max([self.params.dt_max, L.params.dt])

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

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

549

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

551

552

553class LogCFL(Hooks):

554

555 def post_step(self, step, level_number):

556 """

557 Record CFL limit.

558

559 Args:

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

561 level_number (int): the current level number

562

563 Returns:

564 None

565 """

566 super().post_step(step, level_number)

567

568 L = step.levels[level_number]

569

570 self.add_to_stats(

571 process=step.status.slot,

572 time=L.time + L.dt,

573 level=L.level_index,

574 iter=step.status.iter,

575 sweep=L.status.sweep,

576 type='CFL_limit',

577 value=L.status.CFL_limit,

578 )

579

580

581class LogAnalysisVariables(Hooks):

582

583 def post_step(self, step, level_number):

584 """

585 Record Nusselt numbers.

586

587 Args:

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

589 level_number (int): the current level number

590

591 Returns:

592 None

593 """

594 super().post_step(step, level_number)

595

596 L = step.levels[level_number]

597 P = L.prob

598

599 L.sweep.compute_end_point()

600 Nusselt = P.compute_Nusselt_numbers(L.uend)

601 buoyancy_production = P.compute_buoyancy_generation(L.uend)

602 viscous_dissipation = P.compute_viscous_dissipation(L.uend)

603

604 for key, value in zip(

605 ['Nusselt', 'buoyancy_production', 'viscous_dissipation'],

606 [Nusselt, buoyancy_production, viscous_dissipation],

607 ):

608 self.add_to_stats(

609 process=step.status.slot,

610 time=L.time + L.dt,

611 level=L.level_index,

612 iter=step.status.iter,

613 sweep=L.status.sweep,

614 type=key,

615 value=value,

616 )