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