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

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_bottom'], 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 get_fig(self): # pragma: no cover

261 """

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

263

264 Returns

265 -------

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

267 """

268 import matplotlib.pyplot as plt

269 from mpl_toolkits.axes_grid1 import make_axes_locatable

270

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

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

273 self.cax = []

274 divider = make_axes_locatable(axs[0])

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

276 divider2 = make_axes_locatable(axs[1])

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

278 return self.fig

279

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

281 r"""

282 Plot the solution.

283

284 Parameters

285 ----------

286 u : dtype_u

287 Solution to be plotted

288 t : float

289 Time to display at the top of the figure

290 fig : matplotlib.pyplot.figure.Figure

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

292 quantity : (str)

293 quantity you want to plot

294

295 Returns

296 -------

297 None

298 """

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

300 axs = fig.axes

301

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

303

304 if self.spectral_space:

305 u = self.itransform(u)

306

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

308

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

310 axs[i].set_aspect(1)

311 axs[i].set_title(label)

312

313 if t is not None:

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

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

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

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

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

319

320 def compute_vorticity(self, u):

321 if self.spectral_space:

322 u_hat = u.copy()

323 else:

324 u_hat = self.transform(u)

325 Dz = self.Dz

326 Dx = self.Dx

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

328

329 vorticity_hat = self.spectral.u_init_forward

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

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

332

333 def compute_Nusselt_numbers(self, u):

334 """

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

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

337 advection is present.

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

339

340 Args:

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

342

343 Returns:

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

345 """

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

347

348 DzT_hat = self.spectral.u_init_forward

349

350 if self.spectral_space:

351 u_hat = u.copy()

352 else:

353 u_hat = self.transform(u)

354

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

356

357 # compute vT with dealiasing

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

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

360 _me = self.xp.zeros_like(u_pad)

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

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

363

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

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

366

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

368 integral_V = (

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

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

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

372

373 Nusselt_t = self.comm.bcast(

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

375 )

376 Nusselt_b = self.comm.bcast(

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

378 )

379 Nusselt_no_v_t = self.comm.bcast(

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

381 )

382 Nusselt_no_v_b = self.comm.bcast(

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

384 root=0,

385 )

386

387 return {

388 'V': Nusselt_V,

389 't': Nusselt_t,

390 'b': Nusselt_b,

391 't_no_v': Nusselt_no_v_t,

392 'b_no_v': Nusselt_no_v_b,

393 }

394

395 def compute_viscous_dissipation(self, u):

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

397

398 Lap_u_hat = self.spectral.u_init_forward

399

400 if self.spectral_space:

401 u_hat = u.copy()

402 else:

403 u_hat = self.transform(u)

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

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

406 Lap_u = self.itransform(Lap_u_hat)

407

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

409

410 def compute_buoyancy_generation(self, u):

411 if self.spectral_space:

412 u = self.itransform(u)

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

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

415

416

417class CFLLimit(ConvergenceController):

418

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

420 from pySDC.implementations.hooks.log_step_size import LogStepSize

421

422 controller.add_hook(LogCFL)

423 controller.add_hook(LogStepSize)

424

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

426 """

427 Add the embedded error variable to the error function.

428

429 Args:

430 controller (pySDC.Controller): The controller

431 """

432 self.add_status_variable_to_level('CFL_limit')

433

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

435 """

436 Define default parameters here.

437

438 Default parameters are:

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

440 - dt_max (float): maximal step size

441 - dt_min (float): minimal step size

442

443 Args:

444 controller (pySDC.Controller): The controller

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

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

447

448 Returns:

449 (dict): The updated params dictionary

450 """

451 defaults = {

452 "control_order": -50,

453 "dt_max": np.inf,

454 "dt_min": 0,

455 "cfl": 0.4,

456 }

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

458

459 @staticmethod

460 def compute_max_step_size(P, u):

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

462

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

464 cell_wallz[0] = 1

465 cell_wallz[-1] = -1

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

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

468

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

470

471 if P.spectral_space:

472 u = P.itransform(u)

473

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

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

476 max_step_size = min([max_step_size_x, max_step_size_z])

477

478 if hasattr(P, 'comm'):

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

480 return float(max_step_size)

481

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

483 if not CheckConvergence.check_convergence(step):

484 return None

485

486 L = step.levels[0]

487 P = step.levels[0].prob

488

489 L.sweep.compute_end_point()

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

491

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

493

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

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

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

497

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

499

500

501class LogCFL(Hooks):

502

503 def post_step(self, step, level_number):

504 """

505 Record CFL limit.

506

507 Args:

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

509 level_number (int): the current level number

510

511 Returns:

512 None

513 """

514 super().post_step(step, level_number)

515

516 L = step.levels[level_number]

517

518 self.add_to_stats(

519 process=step.status.slot,

520 time=L.time + L.dt,

521 level=L.level_index,

522 iter=step.status.iter,

523 sweep=L.status.sweep,

524 type='CFL_limit',

525 value=L.status.CFL_limit,

526 )

527

528

529class LogAnalysisVariables(Hooks):

530

531 def post_step(self, step, level_number):

532 """

533 Record Nusselt numbers.

534

535 Args:

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

537 level_number (int): the current level number

538

539 Returns:

540 None

541 """

542 super().post_step(step, level_number)

543

544 L = step.levels[level_number]

545 P = L.prob

546

547 L.sweep.compute_end_point()

548 Nusselt = P.compute_Nusselt_numbers(L.uend)

549 buoyancy_production = P.compute_buoyancy_generation(L.uend)

550 viscous_dissipation = P.compute_viscous_dissipation(L.uend)

551

552 for key, value in zip(

553 ['Nusselt', 'buoyancy_production', 'viscous_dissipation'],

554 [Nusselt, buoyancy_production, viscous_dissipation],

555 ):

556 self.add_to_stats(

557 process=step.status.slot,

558 time=L.time + L.dt,

559 level=L.level_index,

560 iter=step.status.iter,

561 sweep=L.status.sweep,

562 type=key,

563 value=value,

564 )