Coverage for pySDC/implementations/problem_classes/RayleighBenard.py: 79%
233 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-08 09:26 +0000
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-08 09:26 +0000
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 )
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)
112 self.X, self.Z = self.get_grid()
113 self.Kx, self.Kz = self.get_wavenumbers()
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()
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)
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)
129 self.Dx = Dx
130 self.Dxx = Dxx
131 self.Dz = S1 @ Dz
132 self.Dzz = S2 @ Dzz
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)
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)
148 # mass matrix
149 M_lhs = {i: {i: U02 @ Id} for i in ['u', 'v', 'T']}
150 self.setup_M(M_lhs)
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}})
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 )
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()
184 self.work_counters['rhs'] = WorkCounter()
186 def eval_f(self, u, *args, **kwargs):
187 f = self.f_init
189 if self.spectral_space:
190 u_hat = u.copy()
191 else:
192 u_hat = self.transform(u)
194 f_impl_hat = self.u_init_forward
196 Dz = self.Dz
197 Dx = self.Dx
199 iu, iv, iT, ip = self.index(['u', 'v', 'T', 'p'])
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)
206 if self.spectral_space:
207 f.impl[:] = f_impl_hat
208 else:
209 f.impl[:] = self.itransform(f_impl_hat).real
211 # -------------------------------------------
212 # treat convection explicitly with dealiasing
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)
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
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])
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
240 self.work_counters['rhs']()
241 return f
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'
249 me = self.spectral.u_init
250 iu, iv, iT, ip = self.index(['u', 'v', 'T', 'p'])
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
258 # perturb slightly
259 rng = self.xp.random.default_rng(seed=seed)
261 noise = self.spectral.u_init
262 noise[iT] = rng.random(size=me[iT].shape)
264 me[iT] += noise[iT].real * noise_level * (self.Z - 1) * (self.Z + 1)
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
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.
279 Args:
280 sol: Some solution that does not need to enforce boundary conditions
282 Returns:
283 Modified version of the solution that satisfies Dirichlet BCs.
284 """
285 ultraspherical = self.spectral.axes[-1]
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,))
292 BC_bottom = ultraspherical.get_BC(x=-1, kind='dirichlet')
293 BC_top = ultraspherical.get_BC(x=1, kind='dirichlet')
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)
303 BC_vals = M_I @ rhs
305 sol_half_hat[i, :, -2:] = BC_vals.T
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,))
312 def get_fig(self): # pragma: no cover
313 """
314 Get a figure suitable to plot the solution of this problem
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
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
332 def plot(self, u, t=None, fig=None, quantity='T'): # pragma: no cover
333 r"""
334 Plot the solution.
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
347 Returns
348 -------
349 None
350 """
351 fig = self.get_fig() if fig is None else fig
352 axs = fig.axes
354 imV = axs[1].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real)
356 if self.spectral_space:
357 u = self.itransform(u)
359 imT = axs[0].pcolormesh(self.X, self.Z, u[self.index(quantity)].real)
361 for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity']):
362 axs[i].set_aspect(1)
363 axs[i].set_title(label)
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])
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'])
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
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.
392 Args:
393 u: The solution you want to compute the Nusselt numbers of
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'])
400 DzT_hat = self.spectral.u_init_forward
402 if self.spectral_space:
403 u_hat = u.copy()
404 else:
405 u_hat = self.transform(u)
407 DzT_hat[iT] = (self.Dz @ u_hat[iT].flatten()).reshape(DzT_hat[iT].shape)
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)
416 nusselt_hat = (vT_hat[0] - DzT_hat[iT]) / self.nx
417 nusselt_no_v_hat = (-DzT_hat[iT]) / self.nx
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)
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 )
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 }
447 def compute_viscous_dissipation(self, u):
448 iu, iv = self.index(['u', 'v'])
450 Lap_u_hat = self.spectral.u_init_forward
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)
460 return abs(u[iu] * Lap_u[iu] + u[iv] * Lap_u[iv])
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])
469class CFLLimit(ConvergenceController):
471 def dependencies(self, controller, *args, **kwargs):
472 from pySDC.implementations.hooks.log_step_size import LogStepSize
474 controller.add_hook(LogCFL)
475 controller.add_hook(LogStepSize)
477 def setup_status_variables(self, controller, **kwargs):
478 """
479 Add the embedded error variable to the error function.
481 Args:
482 controller (pySDC.Controller): The controller
483 """
484 self.add_status_variable_to_level('CFL_limit')
486 def setup(self, controller, params, description, **kwargs):
487 """
488 Define default parameters here.
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
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
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)}
511 @staticmethod
512 def compute_max_step_size(P, u):
513 grid_spacing_x = P.X[1, 0] - P.X[0, 0]
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:]
521 iu, iv = P.index(['u', 'v'])
523 if P.spectral_space:
524 u = P.itransform(u)
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])
530 if hasattr(P, 'comm'):
531 max_step_size = P.comm.allreduce(max_step_size, op=MPI.MIN)
532 return float(max_step_size)
534 def get_new_step_size(self, controller, step, **kwargs):
535 if not CheckConvergence.check_convergence(step):
536 return None
538 L = step.levels[0]
539 P = step.levels[0].prob
541 L.sweep.compute_end_point()
542 max_step_size = self.compute_max_step_size(P, L.uend)
544 L.status.CFL_limit = self.params.cfl * max_step_size
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])
550 self.log(f'dt max: {max_step_size:.2e} -> New step size: {L.status.dt_new:.2e}', step)
553class LogCFL(Hooks):
555 def post_step(self, step, level_number):
556 """
557 Record CFL limit.
559 Args:
560 step (pySDC.Step.step): the current step
561 level_number (int): the current level number
563 Returns:
564 None
565 """
566 super().post_step(step, level_number)
568 L = step.levels[level_number]
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 )
581class LogAnalysisVariables(Hooks):
583 def post_step(self, step, level_number):
584 """
585 Record Nusselt numbers.
587 Args:
588 step (pySDC.Step.step): the current step
589 level_number (int): the current level number
591 Returns:
592 None
593 """
594 super().post_step(step, level_number)
596 L = step.levels[level_number]
597 P = L.prob
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)
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 )