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

11class RayleighBenard(GenericSpectralLinear):

12 """

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

15 The equations we solve are

17 u_x + v_z = 0

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

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

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

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

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

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

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

27 The domain, vertical boundary conditions and pressure gauge are

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

30 T(z=+1) = 0

31 T(z=-1) = 2

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

33 integral over p = 0

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

36 facilitate the Dirichlet BCs.

38 Parameters:

39 Prandl (float): Prandl number

40 Rayleigh (float): Rayleigh number

41 nx (int): Horizontal resolution

42 nz (int): Vertical resolution

43 BCs (dict): Can specify boundary conditions here

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

45 comm (mpi4py.Intracomm): Space communicator

46 """

48 dtype_u = mesh

49 dtype_f = imex_mesh

51 def __init__(

52 self,

53 Prandl=1,

54 Rayleigh=2e6,

55 nx=256,

56 nz=64,

57 BCs=None,

58 dealiasing=3 / 2,

59 comm=None,

60 **kwargs,

61 ):

62 """

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

65 Args:

66 Prandl (float): Prandtl number

67 Rayleigh (float): Rayleigh number

68 nx (int): Resolution in x-direction

69 nz (int): Resolution in z direction

70 BCs (dict): Vertical boundary conditions

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

72 comm (mpi4py.Intracomm): Space communicator

73 """

74 BCs = {} if BCs is None else BCs

75 BCs = {

76 'T_top': 0,

77 'T_bottom': 2,

78 'v_top': 0,

79 'v_bottom': 0,

80 'u_top': 0,

81 'u_bottom': 0,

82 'p_integral': 0,

83 **BCs,

84 }

85 if comm is None:

86 try:

87 from mpi4py import MPI

89 comm = MPI.COMM_WORLD

90 except ModuleNotFoundError:

91 pass

92 self._makeAttributeAndRegister(

93 'Prandl',

94 'Rayleigh',

95 'nx',

96 'nz',

97 'BCs',

98 'dealiasing',

99 'comm',

100 localVars=locals(),

101 readOnly=True,

102 )

103

104 bases = [{'base': 'fft', 'N': nx, 'x0': 0, 'x1': 8}, {'base': 'ultraspherical', 'N': nz}]

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

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

107

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

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

110

111 # construct 2D matrices

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

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

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

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

116 Id = self.get_Id()

117

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

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

120

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

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

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

124

125 self.Dx = Dx

126 self.Dxx = Dxx

127 self.Dz = S1 @ Dz

128 self.Dzz = S2 @ Dzz

129

130 kappa = (Rayleigh * Prandl) ** (-1 / 2.0)

131 nu = (Rayleigh / Prandl) ** (-1 / 2.0)

132

133 # construct operators

134 L_lhs = {

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

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

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

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

139 }

140 self.setup_L(L_lhs)

141

142 # mass matrix

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

144 self.setup_M(M_lhs)

145

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

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

148

149 # BCs

150 self.add_BC(

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

152 )

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

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

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

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

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

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

159 self.add_BC(

160 component='u',

161 equation='u',

162 axis=1,

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

164 x=-1,

165 kind='Dirichlet',

166 line=-1,

167 )

168

169 # eliminate Nyquist mode if needed

170 if nx % 2 == 0:

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

172 for component in self.components:

173 self.add_BC(

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

175 )

176 self.setup_BCs()

177

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

179 f = self.f_init

180

181 if self.spectral_space:

182 u_hat = u.copy()

183 else:

184 u_hat = self.transform(u)

185

186 f_impl_hat = self.u_init_forward

187

188 Dz = self.Dz

189 Dx = self.Dx

190

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

192

193 # evaluate implicit terms

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

195 self._L_T_base = self.base_change @ self.L

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

197

198 if self.spectral_space:

199 f.impl[:] = f_impl_hat

200 else:

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

202

203 # -------------------------------------------

204 # treat convection explicitly with dealiasing

205

206 # start by computing derivatives

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

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

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

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

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

212

213 padding = [self.dealiasing, self.dealiasing]

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

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

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

217

218 fexpl_pad = self.xp.zeros_like(u_pad)

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

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

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

222

223 if self.spectral_space:

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

225 else:

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

227

228 return f

229

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

231 assert t == 0

232 assert (

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

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

235

236 me = self.spectral.u_init

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

238

239 # linear temperature gradient

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

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

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

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

244

245 # perturb slightly

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

247

248 noise = self.spectral.u_init

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

250

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

252

253 if self.spectral_space:

254 me_hat = self.spectral.u_init_forward

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

256 return me_hat

257 else:

258 return me

259

260 def apply_BCs(self, sol):

261 """

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

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

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

265

266 Args:

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

268

269 Returns:

270 Modified version of the solution that satisfies Dirichlet BCs.

271 """

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

273

274 if self.spectral_space:

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

276 else:

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

278

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

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

281

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

283 M_I = np.linalg.inv(M)

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

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

286 i = self.index(component)

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

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

289

290 BC_vals = M_I @ rhs

291

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

293

294 if self.spectral_space:

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

296 else:

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

298

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

300 """

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

302

303 Returns

304 -------

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

306 """

307 import matplotlib.pyplot as plt

308 from mpl_toolkits.axes_grid1 import make_axes_locatable

309

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

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

312 self.cax = []

313 divider = make_axes_locatable(axs[0])

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

315 divider2 = make_axes_locatable(axs[1])

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

317 return self.fig

318

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

320 r"""

321 Plot the solution.

322

323 Parameters

324 ----------

325 u : dtype_u

326 Solution to be plotted

327 t : float

328 Time to display at the top of the figure

329 fig : matplotlib.pyplot.figure.Figure

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

331 quantity : (str)

332 quantity you want to plot

333

334 Returns

335 -------

336 None

337 """

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

339 axs = fig.axes

340

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

342

343 if self.spectral_space:

344 u = self.itransform(u)

345

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

347

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

349 axs[i].set_aspect(1)

350 axs[i].set_title(label)

351

352 if t is not None:

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

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

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

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

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

358

359 def compute_vorticity(self, u):

360 if self.spectral_space:

361 u_hat = u.copy()

362 else:

363 u_hat = self.transform(u)

364 Dz = self.Dz

365 Dx = self.Dx

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

367

368 vorticity_hat = self.spectral.u_init_forward

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

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

371

372 def compute_Nusselt_numbers(self, u):

373 """

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

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

376 advection is present.

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

378

379 Args:

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

381

382 Returns:

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

384 """

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

386

387 DzT_hat = self.spectral.u_init_forward

388

389 if self.spectral_space:

390 u_hat = u.copy()

391 else:

392 u_hat = self.transform(u)

393

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

395

396 # compute vT with dealiasing

397 padding = [self.dealiasing, self.dealiasing]

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

399 _me = self.xp.zeros_like(u_pad)

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

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

402

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

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

405

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

407 integral_V = (

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

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

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

411

412 Nusselt_t = self.comm.bcast(

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

414 )

415 Nusselt_b = self.comm.bcast(

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

417 )

418 Nusselt_no_v_t = self.comm.bcast(

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

420 )

421 Nusselt_no_v_b = self.comm.bcast(

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

423 root=0,

424 )

425

426 return {

427 'V': Nusselt_V,

428 't': Nusselt_t,

429 'b': Nusselt_b,

430 't_no_v': Nusselt_no_v_t,

431 'b_no_v': Nusselt_no_v_b,

432 }

433

434 def compute_viscous_dissipation(self, u):

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

436

437 Lap_u_hat = self.spectral.u_init_forward

438

439 if self.spectral_space:

440 u_hat = u.copy()

441 else:

442 u_hat = self.transform(u)

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

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

445 Lap_u = self.itransform(Lap_u_hat)

446

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

448

449 def compute_buoyancy_generation(self, u):

450 if self.spectral_space:

451 u = self.itransform(u)

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

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

454

455

456class CFLLimit(ConvergenceController):

457

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

459 from pySDC.implementations.hooks.log_step_size import LogStepSize

460

461 controller.add_hook(LogCFL)

462 controller.add_hook(LogStepSize)

463

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

465 """

466 Add the embedded error variable to the error function.

467

468 Args:

469 controller (pySDC.Controller): The controller

470 """

471 self.add_status_variable_to_level('CFL_limit')

472

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

474 """

475 Define default parameters here.

476

477 Default parameters are:

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

479 - dt_max (float): maximal step size

480 - dt_min (float): minimal step size

481

482 Args:

483 controller (pySDC.Controller): The controller

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

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

486

487 Returns:

488 (dict): The updated params dictionary

489 """

490 defaults = {

491 "control_order": -50,

492 "dt_max": np.inf,

493 "dt_min": 0,

494 "cfl": 0.4,

495 }

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

497

498 @staticmethod

499 def compute_max_step_size(P, u):

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

501

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

503 cell_wallz[0] = 1

504 cell_wallz[-1] = -1

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

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

507

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

509

510 if P.spectral_space:

511 u = P.itransform(u)

512

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

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

515 max_step_size = min([max_step_size_x, max_step_size_z])

516

517 if hasattr(P, 'comm'):

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

519 return float(max_step_size)

520

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

522 if not CheckConvergence.check_convergence(step):

523 return None

524

525 L = step.levels[0]

526 P = step.levels[0].prob

527

528 L.sweep.compute_end_point()

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

530

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

532

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

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

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

536

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

538

539

540class LogCFL(Hooks):

541

542 def post_step(self, step, level_number):

543 """

544 Record CFL limit.

545

546 Args:

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

548 level_number (int): the current level number

549

550 Returns:

551 None

552 """

553 super().post_step(step, level_number)

554

555 L = step.levels[level_number]

556

557 self.add_to_stats(

558 process=step.status.slot,

559 time=L.time + L.dt,

560 level=L.level_index,

561 iter=step.status.iter,

562 sweep=L.status.sweep,

563 type='CFL_limit',

564 value=L.status.CFL_limit,

565 )

566

567

568class LogAnalysisVariables(Hooks):

569

570 def post_step(self, step, level_number):

571 """

572 Record Nusselt numbers.

573

574 Args:

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

576 level_number (int): the current level number

577

578 Returns:

579 None

580 """

581 super().post_step(step, level_number)

582

583 L = step.levels[level_number]

584 P = L.prob

585

586 L.sweep.compute_end_point()

587 Nusselt = P.compute_Nusselt_numbers(L.uend)

588 buoyancy_production = P.compute_buoyancy_generation(L.uend)

589 viscous_dissipation = P.compute_viscous_dissipation(L.uend)

590

591 for key, value in zip(

592 ['Nusselt', 'buoyancy_production', 'viscous_dissipation'],

593 [Nusselt, buoyancy_production, viscous_dissipation],

594 ):

595 self.add_to_stats(

596 process=step.status.slot,

597 time=L.time + L.dt,

598 level=L.level_index,

599 iter=step.status.iter,

600 sweep=L.status.sweep,

601 type=key,

602 value=value,

603 )