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({'u': {'u': Dx}, 'v': {'v': Dx}, 'T': {'T': Dx}, 'p': {}})

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

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

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

220

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

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

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

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

225

226 fexpl_pad = self.xp.zeros_like(u_pad)

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

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

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

230

231 if self.spectral_space:

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

233 else:

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

235

236 self.work_counters['rhs']()

237 return f

238

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

240 assert t == 0

241 assert (

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

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

244

245 me = self.spectral.u_init

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

247

248 # linear temperature gradient

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

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

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

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

253

254 # perturb slightly

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

256

257 noise = self.spectral.u_init

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

259

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

261

262 if self.spectral_space:

263 me_hat = self.spectral.u_init_forward

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

265 return me_hat

266 else:

267 return me

268

269 def apply_BCs(self, sol):

270 """

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

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

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

274

275 Args:

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

277

278 Returns:

279 Modified version of the solution that satisfies Dirichlet BCs.

280 """

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

282

283 if self.spectral_space:

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

285 else:

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

287

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

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

290

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

292 M_I = np.linalg.inv(M)

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

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

295 i = self.index(component)

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

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

298

299 BC_vals = M_I @ rhs

300

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

302

303 if self.spectral_space:

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

305 else:

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

307

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

309 """

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

311

312 Returns

313 -------

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

315 """

316 import matplotlib.pyplot as plt

317 from mpl_toolkits.axes_grid1 import make_axes_locatable

318

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

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

321 self.cax = []

322 divider = make_axes_locatable(axs[0])

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

324 divider2 = make_axes_locatable(axs[1])

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

326 return self.fig

327

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

329 r"""

330 Plot the solution.

331

332 Parameters

333 ----------

334 u : dtype_u

335 Solution to be plotted

336 t : float

337 Time to display at the top of the figure

338 fig : matplotlib.pyplot.figure.Figure

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

340 quantity : (str)

341 quantity you want to plot

342

343 Returns

344 -------

345 None

346 """

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

348 axs = fig.axes

349

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

351

352 if self.spectral_space:

353 u = self.itransform(u)

354

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

356

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

358 axs[i].set_aspect(1)

359 axs[i].set_title(label)

360

361 if t is not None:

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

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

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

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

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

367

368 def compute_vorticity(self, u):

369 if self.spectral_space:

370 u_hat = u.copy()

371 else:

372 u_hat = self.transform(u)

373 Dz = self.Dz

374 Dx = self.Dx

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

376

377 vorticity_hat = self.spectral.u_init_forward

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

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

380

381 def compute_Nusselt_numbers(self, u):

382 """

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

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

385 advection is present.

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

387

388 Args:

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

390

391 Returns:

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

393 """

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

395

396 DzT_hat = self.spectral.u_init_forward

397

398 if self.spectral_space:

399 u_hat = u.copy()

400 else:

401 u_hat = self.transform(u)

402

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

404

405 # compute vT with dealiasing

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

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

408 _me = self.xp.zeros_like(u_pad)

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

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

411

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

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

414

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

416 integral_V = (

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

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

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

420

421 Nusselt_t = self.comm.bcast(

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

423 )

424 Nusselt_b = self.comm.bcast(

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

426 )

427 Nusselt_no_v_t = self.comm.bcast(

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

429 )

430 Nusselt_no_v_b = self.comm.bcast(

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

432 root=0,

433 )

434

435 return {

436 'V': Nusselt_V,

437 't': Nusselt_t,

438 'b': Nusselt_b,

439 't_no_v': Nusselt_no_v_t,

440 'b_no_v': Nusselt_no_v_b,

441 }

442

443 def compute_viscous_dissipation(self, u):

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

445

446 Lap_u_hat = self.spectral.u_init_forward

447

448 if self.spectral_space:

449 u_hat = u.copy()

450 else:

451 u_hat = self.transform(u)

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

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

454 Lap_u = self.itransform(Lap_u_hat)

455

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

457

458 def compute_buoyancy_generation(self, u):

459 if self.spectral_space:

460 u = self.itransform(u)

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

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

463

464

465class CFLLimit(ConvergenceController):

466

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

468 from pySDC.implementations.hooks.log_step_size import LogStepSize

469

470 controller.add_hook(LogCFL)

471 controller.add_hook(LogStepSize)

472

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

474 """

475 Add the embedded error variable to the error function.

476

477 Args:

478 controller (pySDC.Controller): The controller

479 """

480 self.add_status_variable_to_level('CFL_limit')

481

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

483 """

484 Define default parameters here.

485

486 Default parameters are:

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

488 - dt_max (float): maximal step size

489 - dt_min (float): minimal step size

490

491 Args:

492 controller (pySDC.Controller): The controller

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

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

495

496 Returns:

497 (dict): The updated params dictionary

498 """

499 defaults = {

500 "control_order": -50,

501 "dt_max": np.inf,

502 "dt_min": 0,

503 "cfl": 0.4,

504 }

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

506

507 @staticmethod

508 def compute_max_step_size(P, u):

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

510

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

512 cell_wallz[0] = 1

513 cell_wallz[-1] = -1

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

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

516

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

518

519 if P.spectral_space:

520 u = P.itransform(u)

521

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

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

524 max_step_size = min([max_step_size_x, max_step_size_z])

525

526 if hasattr(P, 'comm'):

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

528 return float(max_step_size)

529

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

531 if not CheckConvergence.check_convergence(step):

532 return None

533

534 L = step.levels[0]

535 P = step.levels[0].prob

536

537 L.sweep.compute_end_point()

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

539

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

541

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

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

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

545

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

547

548

549class LogCFL(Hooks):

550

551 def post_step(self, step, level_number):

552 """

553 Record CFL limit.

554

555 Args:

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

557 level_number (int): the current level number

558

559 Returns:

560 None

561 """

562 super().post_step(step, level_number)

563

564 L = step.levels[level_number]

565

566 self.add_to_stats(

567 process=step.status.slot,

568 time=L.time + L.dt,

569 level=L.level_index,

570 iter=step.status.iter,

571 sweep=L.status.sweep,

572 type='CFL_limit',

573 value=L.status.CFL_limit,

574 )

575

576

577class LogAnalysisVariables(Hooks):

578

579 def post_step(self, step, level_number):

580 """

581 Record Nusselt numbers.

582

583 Args:

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

585 level_number (int): the current level number

586

587 Returns:

588 None

589 """

590 super().post_step(step, level_number)

591

592 L = step.levels[level_number]

593 P = L.prob

594

595 L.sweep.compute_end_point()

596 Nusselt = P.compute_Nusselt_numbers(L.uend)

597 buoyancy_production = P.compute_buoyancy_generation(L.uend)

598 viscous_dissipation = P.compute_viscous_dissipation(L.uend)

599

600 for key, value in zip(

601 ['Nusselt', 'buoyancy_production', 'viscous_dissipation'],

602 [Nusselt, buoyancy_production, viscous_dissipation],

603 ):

604 self.add_to_stats(

605 process=step.status.slot,

606 time=L.time + L.dt,

607 level=L.level_index,

608 iter=step.status.iter,

609 sweep=L.status.sweep,

610 type=key,

611 value=value,

612 )