Coverage for pySDC/implementations/problem_classes/RayleighBenard3D.py: 98%
223 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +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.core.check_convergence import CheckConvergence
9from pySDC.core.problem import WorkCounter
12class RayleighBenard3D(GenericSpectralLinear):
13 """
14 3D Rayleigh-Benard convection, FFT in x and y and ultraspherical in z, IMEX with the nonlinear advection explicit.
16 Rayleigh-Benard Convection is a variation of incompressible Navier-Stokes.
18 The equations we solve are
20 u_x + v_y + w_z = 0
21 T_t - kappa (T_xx + T_yy + T_zz) = -uT_x - vT_y - wT_z
22 u_t - nu (u_xx + u_yy + u_zz) + p_x = -uu_x - vu_y - wu_z
23 v_t - nu (v_xx + v_yy + v_zz) + p_y = -uv_x - vv_y - wv_z
24 w_t - nu (w_xx + w_yy + w_zz) + p_z - T = -uw_x - vw_y - ww_z
26 with u and v the horizontal velocities (in x and y), w the vertical velocity (in z), T the temperature, p the
27 pressure, indices
28 denoting derivatives, kappa=(Rayleigh * Prandtl)**(-1/2) and nu = (Rayleigh / Prandtl)**(-1/2). Everything on the left
29 hand side, that is the viscous part, the pressure gradient and the buoyancy due to temperature are treated
30 implicitly, while the non-linear convection part on the right hand side is integrated explicitly.
32 The domain, vertical boundary conditions and pressure gauge are
34 Omega = [0, Lx) x [0, Ly] x (0, Lz)
35 T(z=+1) = 0
36 T(z=-1) = Lz
37 u(z=+-1) = v(z=+-1) = w(z=+-1) = 0
38 integral over p = 0
40 The spectral discretization uses FFT horizontally, implying periodic BCs, and an ultraspherical method vertically to
41 facilitate the Dirichlet BCs.
43 Parameters:
44 Prandtl (float): Prandtl number
45 Rayleigh (float): Rayleigh number
46 nx (int): Horizontal resolution
47 nz (int): Vertical resolution
48 BCs (dict): Can specify boundary conditions here
49 dealiasing (float): Dealiasing factor for evaluating the non-linear part
50 comm (mpi4py.Intracomm): Space communicator
51 """
53 dtype_u = mesh
54 dtype_f = imex_mesh
56 def __init__(
57 self,
58 Prandtl=1,
59 Rayleigh=1e6,
60 nx=64,
61 ny=64,
62 nz=32,
63 BCs=None,
64 dealiasing=1.5,
65 comm=None,
66 Lz=1,
67 Lx=4,
68 Ly=4,
69 useGPU=False,
70 **kwargs,
71 ):
72 """
73 Constructor. `kwargs` are forwarded to parent class constructor.
75 Args:
76 Prandtl (float): Prandtl number
77 Rayleigh (float): Rayleigh number
78 nx (int): Resolution in x-direction
79 ny (int): Resolution in y-direction
80 nz (int): Resolution in z direction
81 BCs (dict): Vertical boundary conditions
82 dealiasing (float): Dealiasing for evaluating the non-linear part in real space
83 comm (mpi4py.Intracomm): Space communicator
84 Lz (float): Vertical length of the domain
85 Lx (float): Horizontal length of the domain in x-direction
86 Ly (float): Horizontal length of the domain in y-direction
87 useGPU (bool): Run on GPU or CPU
88 """
89 BCs = {} if BCs is None else BCs
90 BCs = {
91 'T_top': 0,
92 'T_bottom': Lz,
93 'w_top': 0,
94 'w_bottom': 0,
95 'v_top': 0,
96 'v_bottom': 0,
97 'u_top': 0,
98 'u_bottom': 0,
99 'p_integral': 0,
100 **BCs,
101 }
102 if comm is None:
103 try:
104 from mpi4py import MPI
106 comm = MPI.COMM_WORLD
107 except ModuleNotFoundError:
108 pass
109 self._makeAttributeAndRegister(
110 'Prandtl',
111 'Rayleigh',
112 'nx',
113 'ny',
114 'nz',
115 'BCs',
116 'dealiasing',
117 'comm',
118 'Lx',
119 'Ly',
120 'Lz',
121 localVars=locals(),
122 readOnly=True,
123 )
125 bases = [
126 {'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx, 'useFFTW': not useGPU},
127 {'base': 'fft', 'N': ny, 'x0': 0, 'x1': self.Ly, 'useFFTW': not useGPU},
128 {'base': 'ultraspherical', 'N': nz, 'x0': 0, 'x1': self.Lz},
129 ]
130 components = ['u', 'v', 'w', 'T', 'p']
131 super().__init__(bases, components, comm=comm, useGPU=useGPU, **kwargs)
133 self.X, self.Y, self.Z = self.get_grid()
134 self.Kx, self.Ky, self.Kz = self.get_wavenumbers()
136 # construct 3D matrices
137 Dzz = self.get_differentiation_matrix(axes=(2,), p=2)
138 Dz = self.get_differentiation_matrix(axes=(2,))
139 Dy = self.get_differentiation_matrix(axes=(1,))
140 Dyy = self.get_differentiation_matrix(axes=(1,), p=2)
141 Dx = self.get_differentiation_matrix(axes=(0,))
142 Dxx = self.get_differentiation_matrix(axes=(0,), p=2)
143 Id = self.get_Id()
145 S1 = self.get_basis_change_matrix(p_out=0, p_in=1)
146 S2 = self.get_basis_change_matrix(p_out=0, p_in=2)
148 U01 = self.get_basis_change_matrix(p_in=0, p_out=1)
149 U12 = self.get_basis_change_matrix(p_in=1, p_out=2)
150 U02 = self.get_basis_change_matrix(p_in=0, p_out=2)
151 self.eliminate_zeros(S1)
152 self.eliminate_zeros(S2)
153 self.eliminate_zeros(Dz)
154 self.eliminate_zeros(Dzz)
156 self.Dx = Dx
157 self.Dxx = Dxx
158 self.Dy = Dy
159 self.Dyy = Dyy
160 self.Dz = S1 @ Dz
161 self.Dzz = S2 @ Dzz
162 self.S2 = S2
163 self.S1 = S1
165 # compute rescaled Rayleigh number to extract viscosity and thermal diffusivity
166 Ra = Rayleigh / (max([abs(BCs['T_top'] - BCs['T_bottom']), np.finfo(float).eps]) * self.axes[2].L ** 3)
167 self.kappa = (Ra * Prandtl) ** (-1 / 2.0)
168 self.nu = (Ra / Prandtl) ** (-1 / 2.0)
170 # construct operators
171 _D = U02 @ (Dxx + Dyy) + Dzz
172 L_lhs = {}
173 L_lhs['p'] = {'u': U01 @ Dx, 'v': U01 @ Dy, 'w': Dz} # divergence free constraint
174 L_lhs['u'] = {'p': U02 @ Dx, 'u': -self.nu * _D}
175 L_lhs['v'] = {'p': U02 @ Dy, 'v': -self.nu * _D}
176 L_lhs['w'] = {'p': U12 @ Dz, 'w': -self.nu * _D, 'T': -U02 @ Id}
177 L_lhs['T'] = {'T': -self.kappa * _D}
178 self.setup_L(L_lhs)
180 # mass matrix
181 _U02 = U02 @ Id
182 M_lhs = {i: {i: _U02} for i in ['u', 'v', 'w', 'T']}
183 self.setup_M(M_lhs)
185 # BCs
186 self.add_BC(
187 component='p', equation='p', axis=2, v=self.BCs['p_integral'], kind='integral', line=-1, scalar=True
188 )
189 self.add_BC(component='T', equation='T', axis=2, x=-1, v=self.BCs['T_bottom'], kind='Dirichlet', line=-1)
190 self.add_BC(component='T', equation='T', axis=2, x=1, v=self.BCs['T_top'], kind='Dirichlet', line=-2)
191 self.add_BC(component='w', equation='w', axis=2, x=1, v=self.BCs['w_top'], kind='Dirichlet', line=-1)
192 self.add_BC(component='w', equation='w', axis=2, x=-1, v=self.BCs['w_bottom'], kind='Dirichlet', line=-2)
193 self.remove_BC(component='w', equation='w', axis=2, x=-1, kind='Dirichlet', line=-2, scalar=True)
194 for comp in ['u', 'v']:
195 self.add_BC(
196 component=comp, equation=comp, axis=2, v=self.BCs[f'{comp}_top'], x=1, kind='Dirichlet', line=-2
197 )
198 self.add_BC(
199 component=comp,
200 equation=comp,
201 axis=2,
202 v=self.BCs[f'{comp}_bottom'],
203 x=-1,
204 kind='Dirichlet',
205 line=-1,
206 )
208 # eliminate Nyquist mode if needed
209 if nx % 2 == 0:
210 Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index()
211 for component in self.components:
212 self.add_BC(
213 component=component, equation=component, axis=0, kind='Nyquist', line=int(Nyquist_mode_index), v=0
214 )
215 if ny % 2 == 0:
216 Nyquist_mode_index = self.axes[1].get_Nyquist_mode_index()
217 for component in self.components:
218 self.add_BC(
219 component=component, equation=component, axis=1, kind='Nyquist', line=int(Nyquist_mode_index), v=0
220 )
221 self.setup_BCs()
223 self.work_counters['rhs'] = WorkCounter()
225 def eval_f(self, u, *args, **kwargs):
226 """
227 Evaluate the right hand side, split into an implicit and an explicit part.
229 The implicit part is -L u, i.e. diffusion, pressure gradient, buoyancy and, in the pressure line, the negative
230 divergence, converted back to the Chebychev-T basis. The explicit part is the advection -(u d/dx + v d/dy + w
231 d/dz) of u, v, w and T, computed in physical space on a grid padded by the dealiasing factor.
233 Args:
234 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space
235 *args: Not used, the right hand side does not depend on time
236 **kwargs: Not used, the right hand side does not depend on time
238 Returns:
239 dtype_f: The right hand side, in the same space as `u`
240 """
241 f = self.f_init
243 if self.spectral_space:
244 u_hat = u.copy()
245 else:
246 u_hat = self.transform(u)
248 f_impl_hat = self.u_init_forward
250 iu, iv, iw, iT, ip = self.index(['u', 'v', 'w', 'T', 'p'])
251 derivative_indices = [iu, iv, iw, iT]
253 # evaluate implicit terms
254 f_impl_hat = -(self.L @ u_hat.flatten()).reshape(u_hat.shape)
255 for i in derivative_indices:
256 f_impl_hat[i] = (self.S2 @ f_impl_hat[i].flatten()).reshape(f_impl_hat[i].shape)
257 f_impl_hat[ip] = (self.S1 @ f_impl_hat[ip].flatten()).reshape(f_impl_hat[ip].shape)
259 if self.spectral_space:
260 self.xp.copyto(f.impl, f_impl_hat)
261 else:
262 f.impl[:] = self.itransform(f_impl_hat).real
264 # -------------------------------------------
265 # treat convection explicitly with dealiasing
267 # start by computing derivatives
268 padding = (self.dealiasing,) * self.ndim
269 derivatives = []
270 u_hat_flat = [u_hat[i].flatten() for i in derivative_indices]
272 _D_u_hat = self.u_init_forward
273 for D in [self.Dx, self.Dy, self.Dz]:
274 _D_u_hat *= 0
275 for i in derivative_indices:
276 self.xp.copyto(_D_u_hat[i], (D @ u_hat_flat[i]).reshape(_D_u_hat[i].shape))
277 derivatives.append(self.itransform(_D_u_hat, padding=padding).real)
279 u_pad = self.itransform(u_hat, padding=padding).real
281 fexpl_pad = self.xp.zeros_like(u_pad)
282 for i in derivative_indices:
283 for i_vel, iD in zip([iu, iv, iw], range(self.ndim), strict=True):
284 fexpl_pad[i] -= u_pad[i_vel] * derivatives[iD][i]
286 if self.spectral_space:
287 self.xp.copyto(f.expl, self.transform(fexpl_pad, padding=padding))
288 else:
289 f.expl[:] = self.itransform(self.transform(fexpl_pad, padding=padding)).real
291 self.work_counters['rhs']()
292 return f
294 def u_exact(self, t=0, noise_level=1e-3, seed=99):
295 """
296 Initial conditions, which are only available at t=0 and for Lz=1. Velocities and temperature are linear in z
297 between their boundary values, the pressure is zero, and the temperature is perturbed with seeded uniformly
298 distributed noise, multiplied by `noise_level` and z (z - Lz), which vanishes at both plates. The vertical
299 velocity w has to have equal boundary values, since a linear w is not divergence free.
301 Args:
302 t (float): Time, has to be 0
303 noise_level (float): Amplitude of the noise
304 seed (int): Seed for the random number generator
306 Returns:
307 dtype_u: Initial conditions, in spectral space if `spectral_space` is set, else in physical space
308 """
309 assert t == 0
310 assert (
311 self.BCs['w_top'] == self.BCs['w_bottom']
312 ), 'Initial conditions are only implemented for zero vertical velocity gradient'
314 me = self.spectral.u_init
315 iu, iw, iT, ip = self.index(['u', 'w', 'T', 'p'])
317 # linear temperature gradient
318 assert self.Lz == 1
319 for comp in ['T', 'u', 'v', 'w']:
320 a = self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']
321 b = self.BCs[f'{comp}_bottom']
322 me[self.index(comp)] = a * self.Z + b
324 # perturb slightly
325 rng = self.xp.random.default_rng(seed=seed)
327 noise = self.spectral.u_init
328 noise[iT] = rng.random(size=me[iT].shape)
330 me[iT] += noise[iT].real * noise_level * self.Z * (self.Z - self.Lz)
332 if self.spectral_space:
333 me_hat = self.spectral.u_init_forward
334 me_hat[:] = self.transform(me)
335 return me_hat
336 else:
337 return me
339 def compute_Nusselt_numbers(self, u):
340 """
341 Compute the various versions of the Nusselt number. This reflects the type of heat transport.
342 If the Nusselt number is equal to one, it indicates heat transport due to conduction. If it is larger,
343 advection is present.
344 Computing the Nusselt number at various places can be used to check the code.
346 Args:
347 u: The solution you want to compute the Nusselt numbers of
349 Returns:
350 dict: Nusselt number averaged over the entire volume and horizontally averaged at the top and bottom as well
351 as computed from thermal and kinetic dissipation.
352 """
353 iu, iv, iw, iT = self.index(['u', 'v', 'w', 'T'])
354 zAxis = self.spectral.axes[-1]
356 if self.spectral_space:
357 u_hat = u.copy()
358 else:
359 u_hat = self.transform(u)
361 # evaluate derivatives in x, y, and z for u, v, w, and T
362 derivatives = []
363 derivative_indices = [iu, iv, iw, iT]
364 u_hat_flat = [u_hat[i].flatten() for i in derivative_indices]
365 _D_u_hat = self.u_init_forward
366 for D in [self.Dx, self.Dy, self.Dz]:
367 _D_u_hat *= 0
368 for i in derivative_indices:
369 self.xp.copyto(_D_u_hat[i], (D @ u_hat_flat[i]).reshape(_D_u_hat[i].shape))
370 derivatives.append(
371 self.itransform(_D_u_hat).real
372 ) # derivatives[0] contains x derivatives, [1] is y and [2] is z
374 DzT_hat = (self.Dz @ u_hat[iT].flatten()).reshape(u_hat[iT].shape)
376 # compute wT with dealiasing
377 padding = (self.dealiasing,) * self.ndim
378 u_pad = self.itransform(u_hat, padding=padding).real
379 _me = self.xp.zeros_like(u_pad)
380 _me[0] = u_pad[iw] * u_pad[iT]
381 wT_hat = self.transform(_me, padding=padding)[0]
383 # compute Nusselt number
384 nusselt_hat = (wT_hat / self.kappa - DzT_hat) * self.axes[-1].L
386 # compute thermal dissipation
387 thermal_dissipation = self.u_init_physical
388 thermal_dissipation[0, ...] = self.kappa * (
389 derivatives[0][iT] ** 2 + derivatives[1][iT] ** 2 + derivatives[2][iT] ** 2
390 )
391 thermal_dissipation_hat = self.transform(thermal_dissipation)[0]
393 # compute kinetic energy dissipation
394 kinetic_energy_dissipation = self.u_init_physical
395 for i in [iu, iv, iw]:
396 kinetic_energy_dissipation[0, ...] += self.nu * (
397 derivatives[0][i] ** 2 + derivatives[1][i] ** 2 + derivatives[2][i] ** 2
398 )
399 kinetic_energy_dissipation_hat = self.transform(kinetic_energy_dissipation)[0]
401 # get coefficients for evaluation on the boundary and vertical integration
402 zInt = zAxis.get_integration_weights()
403 top = zAxis.get_BC(kind='Dirichlet', x=1)
404 bot = zAxis.get_BC(kind='Dirichlet', x=-1)
406 # compute volume averages
407 integral_V = 0
408 integral_V_th = 0
409 integral_V_kin = 0
410 if self.comm.rank == 0:
412 integral_V = (zInt @ nusselt_hat[0, 0]).real * self.axes[0].L * self.axes[1].L / self.nx / self.ny
413 integral_V_th = (
414 (zInt @ thermal_dissipation_hat[0, 0]).real * self.axes[0].L * self.axes[1].L / self.nx / self.ny
415 )
416 integral_V_kin = (
417 (zInt @ kinetic_energy_dissipation_hat[0, 0]).real * self.axes[0].L * self.axes[1].L / self.nx / self.ny
418 )
420 # communicate result across all tasks
421 Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0)
422 Nusselt_t = self.comm.bcast(self.xp.sum(nusselt_hat.real[0, 0, :] * top, axis=-1) / self.nx / self.ny, root=0)
423 Nusselt_b = self.comm.bcast(self.xp.sum(nusselt_hat.real[0, 0] * bot / self.nx / self.ny, axis=-1), root=0)
424 Nusselt_thermal = self.comm.bcast(1 / self.kappa * integral_V_th / self.spectral.V, root=0)
425 Nusselt_kinetic = self.comm.bcast(1 + 1 / self.kappa * integral_V_kin / self.spectral.V, root=0)
427 return {
428 'V': Nusselt_V,
429 't': Nusselt_t,
430 'b': Nusselt_b,
431 'thermal': Nusselt_thermal,
432 'kinetic': Nusselt_kinetic,
433 }
435 def get_frequency_spectrum(self, u):
436 """
437 Compute the frequency spectrum of the velocities in x and y direction in the horizontal plane for every point in
438 z. If the problem is well resolved, the coefficients will decay quickly with the wave number, and the reverse
439 indicates that the resolution is too low.
441 The returned spectrum has three dimensions. The first is for component (i.e. u or v), the second is for every
442 point in z and the third is the energy in every wave number.
444 Args:
445 u: The solution you want to compute the spectrum of
447 Returns:
448 RayleighBenard3D.xp.ndarray: wave numbers
449 RayleighBenard3D.xp.ndarray: spectrum
450 """
451 xp = self.xp
452 indices = slice(0, 2)
454 # transform the solution to be in frequency space in x and y, but real space in z
455 if self.spectral_space:
456 u_hat = self.itransform(u, axes=(-1,))
457 else:
458 u_hat = self.transform(
459 u,
460 axes=(
461 -3,
462 -2,
463 ),
464 )
465 u_hat = self.spectral.redistribute(u_hat, axis=2, forward_output=False)
467 # compute "energy density" as absolute square of the velocity modes
468 energy = (u_hat[indices] * xp.conjugate(u_hat[indices])).real / (self.axes[0].N ** 2 * self.axes[1].N ** 2)
470 # prepare wave numbers at which to compute the spectrum
471 abs_kx = xp.abs(self.Kx[:, :, 0])
472 abs_ky = xp.abs(self.Ky[:, :, 0])
474 unique_k = xp.unique(xp.append(xp.unique(abs_kx), xp.unique(abs_ky)))
475 n_k = len(unique_k)
477 # compute local spectrum
478 local_spectrum = self.xp.empty(shape=(2, energy.shape[3], n_k))
479 for i, k in zip(range(n_k), unique_k, strict=True):
480 mask = xp.logical_or(abs_kx == k, abs_ky == k)
481 local_spectrum[..., i] = xp.sum(energy[indices, mask, :], axis=1)
483 # assemble global spectrum from local spectra
484 k_all = self.comm.allgather(unique_k)
485 unique_k_all = []
486 for k in k_all:
487 unique_k_all = xp.unique(xp.append(unique_k_all, xp.unique(k)))
488 n_k_all = len(unique_k_all)
490 spectra = self.comm.allgather(local_spectrum)
491 spectrum = self.xp.zeros(shape=(2, self.axes[2].N, n_k_all))
492 for ks, _spectrum in zip(k_all, spectra, strict=True):
493 ks = list(ks)
494 unique_k_all = list(unique_k_all)
495 for k in ks:
496 index_global = unique_k_all.index(k)
497 index_local = ks.index(k)
498 spectrum[..., index_global] += _spectrum[..., index_local]
500 return xp.array(unique_k_all), spectrum
502 def get_vertical_profiles(self, u, components):
503 """
504 Compute horizontally averaged vertical profiles from the horizontal mean mode of each component, transformed to
505 physical space in z and broadcast from rank 0, which holds the mean mode.
507 Args:
508 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space
509 components (list of str): Names of the components you want the profiles of
511 Returns:
512 dict: Profile along z for each component, as xp.ndarray
513 """
514 if self.spectral_space:
515 u_hat = u.copy()
516 else:
517 u_hat = self.transform(u)
519 _u_hat = self.axes[-1].itransform(u_hat, axes=(-1,))
521 avgs = {}
522 for c in components:
523 i = self.index(c)
524 avg = self.xp.ascontiguousarray(_u_hat[i, 0, 0, :].real) / self.axes[0].N / self.axes[1].N
525 self.comm.Bcast(avg, root=0)
526 avgs[c] = avg
528 return avgs